[P] FlowJax - Normalizing flows in JAX

Hello everyone,

Hopefully this is of interest to some of you. For those that don't know, normalising flows can be used as black-box unconditional or conditional distribution approximators, that support both exact sampling and density evaluations. For an excellent review see [https://arxiv.org/abs/1912.02762](https://arxiv.org/abs/1912.02762).

I am developing [flowjax](https://github.com/danielward27/flowjax/tree/main), a Python package for normalising flows, distributions and bijections. It uses Jax for automatic differentiation, and the [equinox](https://github.com/patrick-kidger/equinox) framework built by Patrick Kidger to allow for a familiar object-oriented design.

It includes many powerful …

case definitions easy efficiency hello jax machinelearning plug in simplified transformer tricks