flowjax
A JAX library for distributions, bijections, and normalizing flows implemented as Equinox modules.
What is flowjax?
FlowJAX is a Python library built on JAX and Equinox that provides implementations of probability distributions, bijections, and normalizing flow models. It enables flexible density estimation and sampling while maintaining full compatibility with JAX transformations like automatic differentiation and just-in-time compilation.
Machine learning researchers and practitioners working with probabilistic models, normalizing flows, or density estimation who want JAX-compatible implementations with modern flow architectures.
FlowJAX offers a clean, composable interface where all components are PyTree-compatible Equinox modules, enabling seamless integration with JAX's transformation system while providing state-of-the-art normalizing flow implementations and flexible training methods.
Overview
FlowJAX is a Python library built on JAX and Equinox that provides tools for working with probability distributions, bijections, and normalizing flows. It enables flexible density estimation and sampling while maintaining compatibility with JAX transformations like automatic differentiation and just-in-time compilation.
Key Features
- JAX-Compatible PyTrees — All distributions and bijections are implemented as Equinox modules, making them compatible with JAX transformations.
- State-of-the-Art Flow Models — Includes implementations of coupling flows, masked autoregressive flows, block neural autoregressive flows, planar flows, and triangular spline flows.
- Flexible Training — Provides training scripts for maximum likelihood estimation, variational inference, and contrastive learning for sequential neural posterior estimation.
- First-Class Conditional Support — Offers built-in support for conditional distributions and density estimation.
- Inversion Without Known Inverse — Includes a bisection search algorithm that allows inverting some bijections without a known analytical inverse.
Philosophy
FlowJAX emphasizes composability and JAX compatibility, treating distributions and transformations as first-class PyTree objects that seamlessly integrate with the JAX ecosystem for high-performance machine learning research.
Use Cases
Best For
- Implementing normalizing flow models with JAX transformations
- Density estimation tasks requiring JAX compatibility
- Research on novel flow architectures with automatic differentiation
- Building conditional probability models with flow-based approaches
- Training flow models via maximum likelihood or variational inference
- Creating composable probabilistic models as PyTree objects
Not Ideal For
- Production systems requiring stable, long-term APIs without breaking changes
- Teams not already invested in the JAX and Equinox ecosystem
- Projects needing simple, plug-and-play density estimation without custom model tuning
- Real-time applications where JIT compilation overhead or flow model complexity is prohibitive
Pros & Cons
Pros
All distributions and bijections are PyTree-compatible Equinox modules, enabling seamless use with JAX transformations like automatic differentiation and just-in-time compilation without extra boilerplate.
Implements state-of-the-art normalizing flows such as coupling flows, masked autoregressive flows, and neural spline flows, providing a comprehensive toolkit for research and experimentation.
Includes training scripts for maximum likelihood estimation, variational inference, and contrastive learning, supporting diverse probabilistic modeling tasks as highlighted in the README.
Offers built-in first-class support for conditional distributions, making it straightforward to implement conditional density estimation and sampling for complex models.
Cons
The library is in early development with explicit warnings about potential breaking changes between major releases, making it unreliable for production or long-term projects.
Requires familiarity with JAX and Equinox, which can be a steep barrier for teams not already using these tools, limiting adoption outside specialized ML research circles.
Compared to established alternatives like TensorFlow Probability or Pyro, FlowJAX has a smaller community, fewer pre-trained models, and less extensive documentation or third-party integrations.
Frequently Asked Questions
Related Projects
ALXGoogle Research
OryxProbabilistic reasoning and statistical analysis in TensorFlow
BRAXMassively parallel rigidbody physics simulation on accelerator hardware.
MctxMonte Carlo tree search in JAX
Found a gem we're missing?
Open-Awesome is built by the community, for the community. Submit a project, suggest an awesome list, or help improve the catalog on GitHub.