gensbi.normalizing_flows.bijections.transformers#

Elementwise transformers parameterised per-dimension by a conditioner.

Pure functions of (value, params) — no learnable state of their own.

Classes#

Affine

Elementwise affine transform with log-scale clamping.

RQSpline

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

Functions#

_clamp(a, lo, hi)

Clamp with a straight-through gradient (NumPyro IAF trick).

_inv_softplus(y)

Inverse of softplus: x such that softplus(x) == y (y > 0).

_rqs_apply(z, x_knots, y_knots, derivatives, inverse)

Apply the RQ spline (or its inverse) to a scalar z.

Module Contents#

class gensbi.normalizing_flows.bijections.transformers.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[source]#
clamp_min = -5.0[source]#
num_params = 2[source]#
class gensbi.normalizing_flows.bijections.transformers.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[source]#
min_bin_height = 0.001[source]#
min_bin_width = 0.001[source]#
min_derivative = 0.001[source]#
num_bins = 8[source]#
num_params = 23[source]#
gensbi.normalizing_flows.bijections.transformers._clamp(a, lo, hi)[source]#

Clamp with a straight-through gradient (NumPyro IAF trick).

Parameters:

  • a (jax.Array)

  • lo (float)

  • hi (float)

Return type:

jax.Array

gensbi.normalizing_flows.bijections.transformers._inv_softplus(y)[source]#

Inverse of softplus: x such that softplus(x) == y (y > 0).

Parameters:

y (jax.Array)

Return type:

jax.Array

gensbi.normalizing_flows.bijections.transformers._rqs_apply(z, x_knots, y_knots, derivatives, inverse)[source]#

Apply the RQ spline (or its inverse) to a scalar z.

Returns (out, logderiv) where logderiv = log(dy/dx) evaluated at the relevant x (the same forward derivative is used for both directions; the caller flips its sign for forward). Outside [-B, B] the map is the identity (logderiv 0).

Parameters:

  • z (jax.Array)

  • x_knots (jax.Array)

  • y_knots (jax.Array)

  • derivatives (jax.Array)

  • inverse (bool)