surjectors

Surjection layers for density estimation with normalizing flows

Surjectors is a light-weight library for density estimation using inference and generative surjective normalizing flows, i.e., flows that can reduce or increase dimensionality while retaining exact-likelihood training. For completeness, we also provide implementations of conventional NF-based bijections.

Surjectors makes use of

• Haiku's module system for neural networks,
• Distrax for probability distributions and base bijectors,
• Optax for gradient-based optimization,
• JAX for autodiff and XLA computation.

Why surjections? A conventional normalizing flow is a chain of bijections: every layer is invertible and dimensionality-preserving, and the change of variables uses the Jacobian log-determinant. A surjection relaxes invertibility — it may map to a lower- (inference) or higher-dimensional (generative) space — and replaces the log-determinant with an exact likelihood contribution. This lets a flow reduce dimensionality as part of the model while still being trained by maximum likelihood.

Quickstart

A flow is built from three objects: a base distribution, a transform (a bijector, a surjector, or a Chain of them), and a TransformedDistribution that ties them together and exposes log_prob, sample, sample_and_log_prob and inverse_and_log_prob.

import distrax
import haiku as hk
from jax import numpy as jnp, random as jr
from surjectors import Slice, LULinear, Chain, TransformedDistribution
from surjectors.nn import make_mlp

def decoder_fn(n_dim):
    def _fn(z):
        params = make_mlp([32, 32, n_dim * 2])(z)
        means, log_scales = jnp.split(params, 2, -1)
        return distrax.Independent(distrax.Normal(means, jnp.exp(log_scales)))
    return _fn

@hk.without_apply_rng
@hk.transform
def flow(x):
    base_distribution = distrax.Independent(
        distrax.Normal(jnp.zeros(5), jnp.ones(5)), 1
    )
    transform = Chain([Slice(5, decoder_fn(5)), LULinear(5)])
    pushforward = TransformedDistribution(base_distribution, transform)
    return pushforward.log_prob(x)

x = jr.normal(jr.PRNGKey(1), (1, 10))
params = flow.init(jr.PRNGKey(2), x)
lp = flow.apply(params, x)

Everything runs inside hk.transform, so a TransformedDistribution method can be called directly, as above, when only one method is needed. To expose several methods (e.g. both log_prob and sample) from the same set of parameters, use surjectors.util.as_model, which wraps hk.multi_transform and returns each named method as a typed callable on model.apply.

Examples

Narrative walkthroughs, each self-contained enough to paste into a script or notebook:

• examples/introduction.md — an introduction to Surjectors.
• examples/density_estimation.md — un-/conditional density estimation on two-moons.
• examples/dimensionality_reduction.md — dimensionality reduction with surjections.

Installation

Make sure to have a working JAX installation. Depending on whether you want to use CPU/GPU/TPU, follow these instructions.

To install from PyPI, call:

pip install surjectors

To install the latest GitHub <RELEASE>:

pip install git+https://github.com/dirmeier/surjectors@<RELEASE>

API reference

All layers implement forward_and_likelihood_contribution / inverse_and_likelihood_contribution; bijections additionally expose the Distrax-compatible forward_and_log_det / inverse_and_log_det. Import everything from the top-level surjectors namespace.

surjectors

surjectors.bijectors

surjectors.surjectors

surjectors.nn

surjectors.util

Contributing

Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled good first issue.

In order to contribute:

1. Clone surjectors and install uv from here.

2. Install all dependencies using uv sync --all-groups.

3. Install the Git hooks:

   uv run pre-commit install -t pre-commit -t commit-msg
   

4. Create a new branch locally, e.g. git checkout -b feature/my-new-feature.

5. Implement your contribution and ideally a test case.

6. Check your work (see below).

7. Submit a PR 🙂.

Development commands

The project uses uv for everything (there is no Makefile):

uv sync --all-groups
uv run pytest
uv run ruff check surjectors examples
uv run ruff check --fix surjectors examples
uv run ruff format surjectors examples
uv run mypy surjectors examples
uv run pre-commit run --all-files

Citing Surjectors

If you find our work relevant to your research, please consider citing:

@article{dirmeier2024surjectors,
    author = {Simon Dirmeier},
    title = {Surjectors: surjection layers for density estimation with normalizing flows},
    year = {2024},
    journal = {Journal of Open Source Software},
    publisher = {The Open Journal},
    volume = {9},
    number = {94},
    pages = {6188},
    doi = {10.21105/joss.06188}
}