gensbi.models.maf.made

MADE conditioner with concatenation-based conditioning (flowjax-style).

Conditioning variables are concatenated onto the input and given autoregressive rank -1 (below every data dimension), so every output – including the first – may depend on the condition while the condition depends on nothing. This is the standard conditional-MAF approach (Papamakarios et al. 2017; flowjax). NO cross-feature normalisation (LayerNorm/RMSNorm/GroupNorm): MADE hidden units carry the autoregressive rank, so cross-unit statistics would mix ranks and silently break the flow. See spec §6.

The conditioner is a single cohesive module behind the (x, cond) -> (dim, num_params) interface; alternative conditioning schemes (FiLM, T-NAF, …) may be added later as drop-in conditioners.

Classes

MADE	Autoregressive conditioner: (x, cond) -> params of shape (dim, num_params).
MaskedAutoregressive	MADE conditioner coupled with an elementwise transformer: one MAF layer.

Functions

_rank_vectors(dim, nn_width, num_params, cond_dim)

0-indexed MADE ranks for input, hidden, and output units.

Module Contents

class gensbi.models.maf.made.MADE(dim, cond_dim, num_params, nn_width, nn_depth, rngs, zero_init=True, param_dtype=jnp.float32, dtype=jnp.float32, activation=jax.nn.silu)

Bases: flax.nnx.Module

Autoregressive conditioner: (x, cond) -> params of shape (dim, num_params).

Conditioning is by concatenation: cond is appended to x and given autoregressive rank -1, so every output (incl. dim 0) may depend on it.

Parameters:

• dim (int) – Autoregressive (target) dimension.

• cond_dim (int) – Conditioning dimension; 0 for unconditional.

• num_params (int) – Transform parameters per dimension (e.g. 2 for an Affine transformer).

• nn_width (int) – Width of each masked hidden layer.

• nn_depth (int) – Number of masked hidden layers.

• rngs (nnx.Rngs) – Flax RNG container for parameter initialisation.

• zero_init (bool, optional) – If True (default), zero-initialise the output layer so that all transform parameters start at 0 (Affine becomes the identity).

• param_dtype (DTypeLike, optional) – Dtype for all kernel and bias parameters. Default is float32.

• dtype (DTypeLike, optional) – Compute dtype forwarded to each MaskedLinear. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast. Log-det accumulation in MaskedAutoregressive is unaffected by this knob.

• activation (Callable, optional) – Element-wise activation applied after each hidden layer. Default is jax.nn.silu().

__call__(x, cond=None)

Compute the transform-parameter array from input and optional conditioning.

Parameters:

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

• cond (Array or None, optional) – Conditioning input of shape (cond_dim,), or None for an unconditional conditioner. Required when cond_dim > 0.

Returns:

Transform-parameter array of shape (dim, num_params).

Return type:

Array

Raises:

ValueError – If cond_dim > 0 and cond is None.

activation

cond_dim

dim

hidden_layers

input_layer

num_params

output_layer

class gensbi.models.maf.made.MaskedAutoregressive(dim, cond_dim, transformer, nn_width, nn_depth, rngs, zero_init=True, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: gensbi.normalizing_flows.bijections.base.Bijection

MADE conditioner coupled with an elementwise transformer: one MAF layer.

Implements the Bijection contract. inverse() maps data to noise in a single parallel MADE pass (fast); forward() maps noise to data via a sequential lax.scan over dimensions (slow).

Parameters:

• dim (int) – Dimensionality of the target variable.

• cond_dim (int) – Dimensionality of the conditioning input; 0 for unconditional.

• transformer (Bijection) – Elementwise bijection (e.g. Affine) whose parameters are predicted by the MADE network.

• nn_width (int) – Width of each hidden layer in the MADE network.

• nn_depth (int) – Number of hidden layers in the MADE network.

• rngs (nnx.Rngs) – Flax RNG container for parameter initialisation.

• zero_init (bool, optional) – If True (default), zero-initialise the MADE output layer so that the flow starts as an identity transform.

• param_dtype (DTypeLike, optional) – Dtype for the MADE conditioner’s stored parameters. Default is float32.

• dtype (DTypeLike, optional) – Compute dtype forwarded to the MADE conditioner. Default is float32. Log-det accumulation below stays unconditionally fp32 regardless of this knob.

forward(u, cond=None)

Map noise to data (the sampling direction).

Runs a sequential lax.scan over dimensions: each step calls the MADE network on the partially-built output to obtain parameters for the next dimension. Because dimension i’s parameters depend only on dimensions < i (already final), the per-dimension log-determinant is accumulated inside the scan, avoiding a second full MADE pass.

Parameters:

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

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

Returns:

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

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

inverse(x, cond=None)

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

Runs a single parallel MADE forward pass to obtain the transform parameters, then applies the elementwise transformer inverse to x.

Parameters:

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

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

Returns:

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

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

dim

made

transformer

gensbi.models.maf.made._rank_vectors(dim, nn_width, num_params, cond_dim)

0-indexed MADE ranks for input, hidden, and output units.

With conditioning (cond_dim > 0) the conditioning inputs get rank -1 (before every data dim) and hidden ranks are shifted into [-1, dim-2] so some hidden units carry only the condition and can feed output dim 0.