gensbi.models.maf

Masked Autoregressive Flow (MAF) density model.

Provides MAFlow, a normalizing flow for exact log-density evaluation and sampling built from stacked autoregressive layers, and its configuration dataclass MAFlowParams.

Submodules

• gensbi.models.maf.made
• gensbi.models.maf.masked_linear
• gensbi.models.maf.masks
• gensbi.models.maf.model

Classes

MAFlow	Masked Autoregressive Flow for exact density evaluation and sampling.
MAFlowParams	Architecture parameters for MAFlow.

Package Contents

class gensbi.models.maf.MAFlow(params)

Bases: flax.nnx.Module

Masked Autoregressive Flow for exact density evaluation and sampling.

Stacks MaskedAutoregressive layers separated by permutations, with an optional data-end Standardize bijection, over a standard-normal base distribution.

Log-density follows the change-of-variables formula: log_prob(x, cond) = base.log_prob(u) + logdet, where u, logdet = chain.inverse(x, cond). The base distribution is built lazily and never enters nnx state.

Parameters:

params (MAFlowParams) – Full architecture configuration; see MAFlowParams.

_base()

log_prob(x, cond=None)

Compute the change-of-variables log-density for a batch of samples.

Parameters:

• x (Array) – Data batch. Shape (B, dim) when channels == 1, or (B, dim, C) when channels > 1 (the channel axis is flattened internally to (B, dim * C)).

• cond (Array or None, optional) – Conditioning batch of shape (B, cond_dim) for cond_channels == 1, or (B, cond_dim, C_cond) for cond_channels > 1 (also flattened internally). Pass None for an unconditional model.

Returns:

Log-probability of shape (B,).

Return type:

Array

sample(key, cond=None, nsamples=None)

Draw samples from the flow.

Parameters:

• key (jax.random.PRNGKey) – Random key.

• cond (Array or None, optional) – Conditioning batch of shape (nsamples, cond_dim) for cond_channels == 1, or (nsamples, cond_dim, C_cond) for cond_channels > 1 (flattened internally). If provided, the number of samples is inferred from cond.shape[0] and nsamples is ignored.

• nsamples (int or None, optional) – Number of samples to draw. Required when cond is None.

Returns:

Sample array of shape (nsamples, dim, channels) for all C >= 1 (C = 1 gives (nsamples, dim, 1); channel axis is never collapsed).

Return type:

Array

set_standardization(mean, std)

Set the data-end Standardize bijection’s mean/std buffers in place.

Accepts shapes (dim,) (broadcast to (dim, 1)), (dim, 1), (C,) (per-channel broadcast), or a scalar broadcastable to (dim, channels).

Raises ValueError if built with standardize=False.

Return type:

None

chain

channels

cond_channels

cond_dim

dim

flat_dim

class gensbi.models.maf.MAFlowParams

Architecture parameters for MAFlow.

Only rngs and dim are required. transformer defaults to Affine() (pass RQSpline() for a spline flow).

Parameters:

• rngs (nnx.Rngs) – Flax RNG container used to initialise all trainable parameters.

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

• cond_dim (int, optional) – Dimensionality of the conditioning input. Default is 0 (unconditional).

• n_layers (int, optional) – Number of MaskedAutoregressive layers. Default is 5.

• transformer (Bijection or None, optional) – Elementwise bijection used by each autoregressive layer. If None (default), an Affine bijection is constructed automatically in __post_init__.

• nn_width (int, optional) – Width of each hidden layer in the MADE conditioner network. Default is 64.

• nn_depth (int, optional) – Number of hidden layers in the MADE conditioner network. Default is 2.

• permutation (str, optional) – Permutation applied between autoregressive layers. "reverse" (default) reverses the dimension ordering; "random" applies a random shuffle sampled at construction time.

• standardize (bool, optional) – If True (default), append a Standardize bijection at the data end of the chain.

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

• param_dtype (DTypeLike, optional) – Dtype for all stored (master) MADE kernel/bias parameters. Default is float32.

• dtype (DTypeLike, optional) – Compute dtype knob threaded through the MADE conditioners. Default is float32 (unlike the bf16-default DiT-family models, MAF keeps fp32 compute by default pending dedicated stability testing — see the mixed-precision design spec). Log-det accumulation is unconditionally fp32 regardless of this knob.

__post_init__()

channels: int = 1

cond_channels: int = 1

cond_dim: int = 0

dim: int

dtype: jax.typing.DTypeLike

n_layers: int = 5

nn_depth: int = 2

nn_width: int = 64

param_dtype: jax.typing.DTypeLike

permutation: str = 'reverse'

rngs: flax.nnx.Rngs

standardize: bool = True

transformer: gensbi.normalizing_flows.bijections.base.Bijection | None = None

zero_init: bool = True