jax-nf TFP/JAX implementation of Normalizing Flows The idea is just to implement a few things in order to be able to write NFs with a combination of Flax, Jax, and TFP.