Developing a Python package with standard normalizing flow implementations that is based on JAX and Flax.
JAX
Flax