FlowJax

FlowJax is a package for creating and fitting normalising flows in Jax.

Installation

pip install flowjax

Examples

• Unconditional density estmation
• Conditional density estimation
• Variational inference
• Bounded flows
• Sequential neural posterior estimation

API

• Bijections
• Distributions
• Experimental
• Flows
• Loss functions
• Training
• Wrappers

Miscellaneous

• FAQ