flowjax (pypi): Normalizing flow implementations in jax.
Keywords: jax|normalising-flows|normalizing-flows