FlowJAX

FlowJAX is a package for continuous distributions, bijections, and normalizing flows. It is built on top of equinox and jax.

Key features include:

• Wide range of distributions and bijections: Provides many standard and specialized mathematical objects.
• JAX Compatibility: Distributions and bijections are implemented as PyTrees via equinox modules, allowing them to work seamlessly with JAX transformations like jit, grad, and vmap.
• Normalizing Flows: Includes many state-of-the-art normalizing flow models.
• Conditional Support: First-class support for conditional distributions, which is essential for applications like amortized variational inference and simulation-based inference.

In FlowJAX, normalizing flows are implemented as convenient constructors for flowjax.distributions.Transformed distributions. A flow defines a flexible distribution by transforming a simple base distribution through a flexible bijection (diffeomorphism).

FlowJAX uses two primary methods to compose bijections:

1. flowjax.bijections.Chain: Allows chaining arbitrary and heterogeneous bijections. However, it compiles each layer separately, which can increase compilation time.
2. flowjax.bijections.Scan: Requires layers to share the same structure, but avoids compiling each layer separately.

Note: All FlowJAX flows use flowjax.bijections.Scan by default to optimize and reduce compilation times.

Conditional distributions in FlowJAX include a cond_shape attribute (which is None for unconditional distributions). When using conditional distributions, you can pass a condition argument to sample and log_prob. If you pass batches of samples and conditions, FlowJAX follows NumPy's broadcasting rules.

from flowjax.distributions import Normal
from flowjax.flows import coupling_flow
import jax.numpy as jnp
import jax.random as jr

key = jr.key(0)
dist = coupling_flow(key, base_dist=Normal(jnp.zeros(3)), cond_dim=2)

# Sampling 10 times for a single conditioning variable instance
condition = jnp.ones(2)
samples = dist.sample(key, (10,), condition=condition)
# samples.shape == (10, 3)

# Sampling once for each of 5 conditioning variable instances
condition = jnp.ones((5, 2))
samples = dist.sample(key, condition=condition)
# samples.shape == (5, 3)

Distributions and bijections are registered as PyTrees via equinox. This allows you to use equinox.filter_vmap to vectorize the initialization and operations of distributions. Instead of managing a batch_shape attribute manually, you can use JAX vectorization to handle batches of non-identically distributed objects.

import equinox as eqx
import jax.numpy as jnp
from flowjax.distributions import Normal

# Vectorize initialization to create a batch of independent distributions
normals = eqx.filter_vmap(Normal)(jnp.arange(3))

# Vectorize log_prob computation over the batch of distributions
log_probs = eqx.filter_vmap(lambda dist, x: dist.log_prob(x))(normals, jnp.arange(3))
# log_probs.shape == (3,)

Bijections in normalizing flows often have asymmetric computational efficiency. In FlowJAX, the default behavior is optimized for density evaluation (log_prob) rather than sampling.

• Default Behavior: FlowJAX inverts the bijection with flowjax.bijections.Invert before transforming the base distribution. This makes log_prob faster but makes sample and sample_and_log_prob slower.
• Optimization for Sampling: If you require faster sample and sample_and_log_prob methods (for example, when using certain variational objectives), you should set invert=False when constructing the flow.

To build the HTML version of the FlowJAX documentation on your local machine, navigate to the docs directory and run the make command. Once the build is complete, you can view the documentation by opening the generated index.html file in your browser.

make -C docs html
# View the output at ./docs/_build/html/index.html

When working with target samples that require specific scales or forms, you can use a bijection for preprocessing. To ensure the flow operates on the original scale, wrap the flow in a Transformed distribution using the inverse of your preprocessing bijection.

import jax
import jax.numpy as jnp
import jax.random as jr
from flowjax.distributions import Normal, Transformed
from flowjax.bijections import Affine, Invert
from flowjax.train import fit_to_data

key = jr.key(0)
x = jr.normal(key, (1000, 3))
flow = Normal(jnp.ones(3))

# 1. Define preprocessing
preprocess = Affine(-x.mean(axis=0)/x.std(axis=0), 1/x.std(axis=0))
x_processed = jax.vmap(preprocess.transform)(x)

# 2. Fit the flow to processed data
# flow, losses = fit_to_data(key, dist=flow, x=x_processed)

# 3. Wrap the flow to 'undo' preprocessing automatically
flow = Transformed(flow, Invert(preprocess))

To enable helpful runtime error messages regarding array shapes and types, you can use jaxtyping's import hook combined with beartype. This is useful for catching shape mismatches during development.

from jaxtyping import install_import_hook

# Use the hook to wrap the flowjax import with beartype
with install_import_hook("flowjax", "beartype.beartype"):
    from flowjax import bijections as bij

# This will now raise a helpful error if shapes are incorrect
exp = bij.Exp(shape=2)  # Raises error because shape should be a tuple, e.g., (2,)