Matt Wescott

This variation on @tomhennigan's [example](https://github.com/deepmind/dm-haiku/issues/32#issuecomment-619519490) tries to build a tree of module types. It assumes the parameter creation order is preserved when flattening the parameter dictionary, which may be incorrect....

Autoregressive flow on jax

5

The docstring for `MaskedAutoregressiveFlow` suggests `AutoregressiveNetwork`, which fails on the jax substrate due to Keras dependencies. If there are still missing pieces, are you interested in contributions? @sharadmv @SiegeLordEx