Matt Wescott

@tomhennigan your `get_parameter_tree` proposal would be useful to me. On a related note, for some parameter transformations it is useful to know the type of the corresponding module. Is this...

@tomhennigan thanks for the examples. > I could imagine extending this custom getter to also pass the module as well as the init function This would be great, so much...

@tomhennigan Would it be impractical to instead intercept module creation? With a mapping from module names to types, could use `tree.flatten_with_path_up_to` for straightforward type-dependent transformations of the parameter tree. Either...

Thanks @brianwa84 and @SiegeLordEx. [Here](https://github.com/mattwescott/probability/blob/d2f97d70efc4b36d41cf08848d8ad7f8e881cd5a/discussion/examples/autoregressive_flow_on_jax.ipynb) is an example with multiple layers using Haiku. This approach seems tolerable for my current usecases, but I could be missing something important. Comments, ideas,...