gensbi.normalizing_flows

Pure normalizing-flow abstractions (bijections + change-of-variables).

Concrete flow models live under gensbi.models (MAFlow, TarFlow).

Submodules

• gensbi.normalizing_flows.bijections

Classes

Affine	Elementwise affine transform with log-scale clamping.
Bijection	Invertible map between a noise variable and a data variable.
Chain	Sequential composition of bijections.
Mask	Fixed (non-trainable) buffer for bijection metadata.
Permutation	Reorder dimensions; conditioning is ignored; log-det is always 0.
RQSpline	Elementwise monotonic rational-quadratic spline on [-B, B].
Standardize	Fixed affine standardization using non-trainable mean and std buffers.

Package Contents

class gensbi.normalizing_flows.Affine(clamp_min=-5.0, clamp_max=3.0)

Elementwise affine transform with log-scale clamping.

The parameter layout per dimension is [shift mu, log-scale a] (num_params == 2). Log-scale values are clamped to [clamp_min, clamp_max] before exponentiation using a straight-through gradient (NumPyro IAF trick).

Forward (noise→data): x = u * exp(a) + mu; logabsdet = +sum(a). Inverse (data→noise): u = (x - mu) * exp(-a); logabsdet = -sum(a).

Parameters:

• clamp_min (float, optional) – Lower bound for log-scale clamping. Default is -5.0.

• clamp_max (float, optional) – Upper bound for log-scale clamping. Default is 3.0.

_split(params)

Parameters:

params (jax.Array)

Return type:

tuple[jax.Array, jax.Array]

forward(u, params)

Map noise to data: x = u * exp(a) + mu.

Parameters:

• u (Array) – Noise-space input of shape (dim,).

• params (Array) – Per-dimension parameters of shape (dim, 2); column 0 is the shift mu and column 1 is the log-scale a.

Returns:

• x (Array) – Data-space output.

• logabsdet (Array) – Sum of clamped log-scales: sum(a).

Return type:

tuple[jax.Array, jax.Array]

forward_dim(u_i, params_i)

Scalar noise-to-data transform for a single dimension.

Used by the sequential sampling scan in the autoregressive loop, which accumulates the per-dimension log-determinant as it goes.

Parameters:

• u_i (Array) – Scalar noise-space value for one dimension.

• params_i (Array) – Parameter vector of length 2 for the same dimension; index 0 is the shift mu, index 1 is the log-scale a.

Returns:

• x_i (Array) – Scalar data-space output: u_i * exp(a) + mu.

• logabsdet_i (Array) – Forward log-determinant contribution for this dimension: a.

Return type:

tuple[jax.Array, jax.Array]

inverse(x, params)

Map data to noise: u = (x - mu) * exp(-a).

Parameters:

• x (Array) – Data-space input of shape (dim,).

• params (Array) – Per-dimension parameters of shape (dim, 2); column 0 is the shift mu and column 1 is the log-scale a.

Returns:

• u (Array) – Noise-space output.

• logabsdet (Array) – Negative sum of clamped log-scales: -sum(a).

Return type:

tuple[jax.Array, jax.Array]

clamp_max = 3.0

clamp_min = -5.0

num_params = 2

class gensbi.normalizing_flows.Bijection

Bases: flax.nnx.Module

Invertible map between a noise variable and a data variable.

Subclasses implement a differentiable bijection with a direction convention fixed across the library: forward() maps noise to data (the sampling direction) and inverse() maps data to noise (the density-evaluation direction). Both directions also return the log absolute determinant of the Jacobian of the transform they apply.

abstractmethod forward(u, cond=None)

Map noise to data (the sampling direction).

Parameters:

• u (Array) – Noise-space input.

• cond (Array or None, optional) – Conditioning input, or None for an unconditional map.

Returns:

• x (Array) – Data-space output.

• logabsdet (Array) – Log absolute determinant of the Jacobian of the forward map.

Raises:

NotImplementedError – This is an abstract method; subclasses must override it.

Return type:

tuple[jax.Array, jax.Array]

abstractmethod inverse(x, cond=None)

Map data to noise (the density-evaluation direction).

Parameters:

• x (Array) – Data-space input.

• cond (Array or None, optional) – Conditioning input, or None for an unconditional map.

Returns:

• u (Array) – Noise-space output.

• logabsdet (Array) – Log absolute determinant of the Jacobian of the inverse map.

Raises:

NotImplementedError – This is an abstract method; subclasses must override it.

Return type:

tuple[jax.Array, jax.Array]

class gensbi.normalizing_flows.Chain(bijections)

Bases: gensbi.normalizing_flows.bijections.base.Bijection

Sequential composition of bijections.

Applies bijections in the order given for forward() (noise→data) and in reversed order for inverse() (data→noise). Log-absolute determinants accumulate by summation. bijections[-1] is closest to data space; it is the last bijection applied in forward() and the first applied in inverse().

Parameters:

bijections (list of Bijection) – Ordered list of bijections stored in noise→data (forward) order.

forward(u, cond=None)

Map noise to data by composing all bijections in order.

Parameters:

• u (Array) – Noise-space input.

• cond (Array or None, optional) – Conditioning input passed to every bijection, or None.

Returns:

• x (Array) – Data-space output.

• logabsdet (Array) – Accumulated log absolute determinant of the composed forward map.

inverse(x, cond=None)

Map data to noise by composing all bijections in reversed order.

Parameters:

• x (Array) – Data-space input.

• cond (Array or None, optional) – Conditioning input passed to every bijection, or None.

Returns:

• u (Array) – Noise-space output.

• logabsdet (Array) – Accumulated log absolute determinant of the composed inverse map.

bijections

class gensbi.normalizing_flows.Mask(value, *, hijax=None, ref=None, eager_sharding=None, **metadata)

Bases: flax.nnx.Variable

Fixed (non-trainable) buffer for bijection metadata.

A non-Param Variable subclass so that optimizers and EMA utilities skip it while checkpointing still saves and restores it. Typical use: autoregressive masks in MaskedAutoregressive.

Parameters:

• value (A | VariableMetadata[A])

• hijax (bool | None)

• ref (bool | None)

• eager_sharding (bool | None)

• metadata (Any)

class gensbi.normalizing_flows.Permutation(perm)

Bases: gensbi.normalizing_flows.bijections.base.Bijection

Reorder dimensions; conditioning is ignored; log-det is always 0.

Both the permutation index array and its inverse are stored as Mask buffers so that optimizers and EMA utilities skip them.

Parameters:

perm (Array) – Integer index array of shape (dim,) that defines the reordering. perm[i] is the source index for output position i.

forward(u, cond=None)

Map noise to data by applying the inverse permutation.

Parameters:

• u (Array) – Noise-space input of shape (dim,).

• cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

• x (Array) – Reordered data-space output.

• logabsdet (Array) – Zero scalar (permutations have unit Jacobian determinant).

inverse(x, cond=None)

Map data to noise by applying perm to reorder dimensions.

Parameters:

• x (Array) – Data-space input of shape (dim,).

• cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

• u (Array) – Reordered noise-space output.

• logabsdet (Array) – Zero scalar (permutations have unit Jacobian determinant).

classmethod random(dim, rngs)

Construct a uniformly random permutation.

Parameters:

• dim (int) – Number of dimensions.

• rngs (nnx.Rngs) – Flax RNG container used to draw the random permutation index.

Returns:

A Permutation with a randomly shuffled index array.

Return type:

Permutation

classmethod reverse(dim)

Construct a permutation that reverses dimension order.

Parameters:

dim (int) – Number of dimensions.

Returns:

A Permutation that reverses the dimension order.

Return type:

Permutation

inv_perm

perm

class gensbi.normalizing_flows.RQSpline(num_bins=8, range_bound=5.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)

Elementwise monotonic rational-quadratic spline on [-B, B].

Linear (identity) tails outside the spline interval. The mapping is parameterised per dimension; with zero-initialised parameters the spline reduces to the identity, so the flow warm-starts as a standard normal (same warm-start contract as Affine).

The parameter layout per dimension is [widths(K), heights(K), inner_derivatives(K-1)] of total length 3K - 1. Reference: Durkan et al. 2019 (https://arxiv.org/abs/1906.04032).

Parameters:

• num_bins (int, optional) – Number of spline bins K. Default is 8.

• range_bound (float, optional) – Half-width B of the spline interval [-B, B]. Default is 5.0.

• min_bin_width (float, optional) – Minimum fractional bin width after softmax normalisation. Default is 1e-3.

• min_bin_height (float, optional) – Minimum fractional bin height after softmax normalisation. Default is 1e-3.

• min_derivative (float, optional) – Minimum knot derivative value. Default is 1e-3.

_fwd_scalar(x, params)

Parameters:

• x (jax.Array)

• params (jax.Array)

_inv_scalar(u, params)

Parameters:

• u (jax.Array)

• params (jax.Array)

_knots(params)

Raw params -> (x_knots, y_knots, derivatives), each over K+1 knots.

Parameters:

params (jax.Array)

forward(u, params)

Map noise to data (vectorised inverse spline pass).

Parameters:

• u (Array) – Noise-space input of shape (dim,).

• params (Array) – Per-dimension spline parameters of shape (dim, 3K-1).

Returns:

• x (Array) – Data-space output.

• logabsdet (Array) – Negative sum of log-derivatives of the spline: -sum(log g'(x)), where g' is the forward spline derivative evaluated at the corresponding data-space position.

forward_dim(u_i, params_i)

Scalar noise-to-data transform for a single dimension.

Used by the sequential sampling scan in the autoregressive loop, which accumulates the per-dimension log-determinant as it goes.

Parameters:

• u_i (Array) – Scalar noise-space value for one dimension.

• params_i (Array) – Spline parameter vector of length 3K-1 for the same dimension.

Returns:

• x_i (Array) – Scalar data-space output.

• logabsdet_i (Array) – Forward log-determinant contribution for this dimension: -log g'(x_i) (matching forward()).

Return type:

tuple[jax.Array, jax.Array]

inverse(x, params)

Map data to noise (fast vectorised spline pass).

Parameters:

• x (Array) – Data-space input of shape (dim,).

• params (Array) – Per-dimension spline parameters of shape (dim, 3K-1).

Returns:

• u (Array) – Noise-space output.

• logabsdet (Array) – Positive sum of log-derivatives of the spline at each input: +sum(log g'(x)), where g' is the forward spline derivative (data→noise direction).

B = 5.0

min_bin_height = 0.001

min_bin_width = 0.001

min_derivative = 0.001

num_bins = 8

num_params = 23

class gensbi.normalizing_flows.Standardize(dim)

Bases: gensbi.normalizing_flows.bijections.base.Bijection

Fixed affine standardization using non-trainable mean and std buffers.

Buffers default to identity (mean 0, std 1) and can be updated in place via set_stats(). They are stored as Mask variables so that optimizers and EMA utilities skip them.

Parameters:

dim (int) – Dimension of the data vector (length of mean and std buffers).

forward(u, cond=None)

Map noise to data by destandardizing: x = u * std + mean.

Parameters:

• u (Array) – Noise-space input of shape (dim,).

• cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

• x (Array) – Destandardized data-space output.

• logabsdet (Array) – Log absolute determinant of the forward map: sum(log std).

inverse(x, cond=None)

Map data to noise by standardizing: u = (x - mean) / std.

Parameters:

• x (Array) – Data-space input of shape (dim,).

• cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

• u (Array) – Standardized noise-space output.

• logabsdet (Array) – Log absolute determinant of the inverse map: -sum(log std).

set_stats(mean, std)

Update the mean and standard-deviation buffers in place.

Parameters:

• mean (Array) – New mean values of shape (dim,).

• std (Array) – New standard-deviation values of shape (dim,); must be strictly positive.

Returns:

This method modifies the buffers in place and returns nothing.

Return type:

None

mean

std