Feature Request: implement conditional flows #269
Description
Distrax is a nicely designed library and I would love to use it, but it does not support conditional flows.
Generalizing the library to support conditional flows should not be difficult, this author has achieved this by subclassing the classes in Distrax and overriding a few methods: https://github.com/smsharma/jax-conditional-flows
Most of the code just deals with adding a new context argument to the existing interfaces and passing that down towards the elemental building blocks.
Maintaining such a code is difficult, however, considering how fast the JAX ecosystem evolves currently.
It would be great if conditional flows were first-class citizen in Distrax. This would enlarge the phase-space of applications considerably.