fleqx

Normalizing flows for JAX, built directly on distreqx.

A flow from fleqx is a plain distreqx.distributions.Transformed — there's no flow-specific wrapper, so log_prob, sample, optax training, etc. all work exactly as they would for any other distreqx distribution. As with any distreqx distribution, batches are handled with jax.vmap rather than by passing in arrays with a leading batch axis.

• fleqx.train.fit accepts an arbitrary pytree for both dist and data -- e.g. a dict of independently-trained distributions and their corresponding data arrays -- not just a single distribution and a single (n, dim) array. Pytree leaves may be None for fields you're not using.
• Each flow constructor also accepts template= in place of dim, for a distribution over an arbitrary pytree of arrays (e.g. a dict of named variables) instead of a flat vector.

Three flow types are implemented so far: coupling, masked autoregressive, and planar.

Example

import jax.numpy as jnp
import jax.random as jr

import fleqx

key = jr.key(0)
flow = fleqx.coupling_flow(key, dim=2)

sample = flow.sample(jr.key(1))
log_p = flow.log_prob(sample)

data = jr.normal(jr.key(2), (1000, 2))
flow, losses = fleqx.fit(jr.key(3), flow, data)

coupling_flow, masked_autoregressive_flow, planar_flow and fit are also available at their fully-qualified paths (fleqx.flows.coupling_flow, fleqx.train.fit, etc.), alongside the rest of each submodule.

Installation

pip install fleqx

That is enough on its own: fleqx takes its distreqx bijectors and distributions through parax, which works with any version of distreqx and supplies anything the installed one is missing.

See the documentation for the full API.

Acknowledgements

Built on distreqx (Owen Lockwood) and parax (freezing/unwrapping of fixed sub-components, e.g. data='s standardizing layer). The bijectors were ported from flowjax (Daniel Ward), a more complete flows library that's worth using directly if you don't specifically need the distreqx API.

Development

pip install -e ".[test]"
pytest

This library was written by Claude, porting the bijectors directly from flowjax. Behaviour should be nearly identical to flowjax's, though minor differences may remain, and the code hasn't yet had a full human review.