gensbi.normalizing_flows.bijections#

Pure normalizing-flow bijection abstractions (Tier 1).

Submodules#

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.bijections.Affine(clamp_min=-5.0, clamp_max=3.0)[source]#

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)[source]#

Parameters:: params (jax.Array)
Return type:: tuple[jax.Array, jax.Array]

forward(u, params)[source]#

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)[source]#

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)[source]#

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.bijections.Bijection[source]#

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)[source]#

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)[source]#

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.bijections.Chain(bijections)[source]#

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)[source]#

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)[source]#

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.bijections.Mask(value, *, hijax=None, ref=None, eager_sharding=None, **metadata)[source]#

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.bijections.Permutation(perm)[source]#

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)[source]#

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)[source]#

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)[source]#

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)[source]#

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.bijections.RQSpline(num_bins=8, range_bound=5.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)[source]#

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)[source]#

Parameters:

x (jax.Array)
params (jax.Array)

_inv_scalar(u, params)[source]#

Parameters:

u (jax.Array)
params (jax.Array)

_knots(params)[source]#

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

Parameters:: params (jax.Array)

forward(u, params)[source]#

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)[source]#

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)[source]#

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.bijections.Standardize(dim)[source]#

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)[source]#

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)[source]#

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)[source]#

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#