NuX

A JAX library for normalizing flows with bijective, dimension-changing and surjective layers.

NuX is a library for building, training and evaluating normalizing flows in JAX. I wrote it as the Information Fusion Lab’s flow library during my PhD, and it has 87 GitHub stars. It was the code base for my early research on flows that change dimension, including the rectangular change of variables formula from AISTATS 2021.

Key features

  • Bijective layers including affine and nonlinear couplings, rational-quadratic splines, residual flows, convex potential flows and FFJORD
  • Dimension-changing layers including injective slices and the tall and wide matrix-vector products of the rectangular change of variables
  • Surjective layers for dequantization, padding and max pooling
  • Priors including Gaussian, mixture and discrete distributions
  • Data-dependent initialization, so each layer sets its parameters from the first batch it sees

Quick start

import jax
import jax.numpy as jnp
import nux

rng_key = jax.random.PRNGKey(0)
x = jax.random.normal(rng_key, (8, 3))

flow = nux.Sequential([nux.DenseMVP(),
                       nux.SneakyReLU(),
                       nux.UnitGaussianPrior()])

# The first call initializes the parameters from the data
z, log_px = flow(x, rng_key=rng_key)
params = flow.get_params()

# Sample by running the flow in reverse
x_samples, log_px_samples = flow(jnp.zeros_like(x), params=params, inverse=True, rng_key=rng_key)