gensbi.recipes.flow_pipeline#

NPE pipeline for discrete normalizing flows (parallel track).

The flow IS the density model: no ConditionalWrapper, no GenerativeMethod. Trains q(obs | cond) by max-likelihood. NPE convention: obs = theta, cond = x (mirrors ConditionalPipeline so the diagnostics run unchanged).

Classes#

ConditionalFlowPipeline

Max-likelihood NPE pipeline wrapping an MAFlow.

Functions#

_warn_unused_kwargs(kwargs)

Warn that solver-style kwargs are ignored by the (solver-free) flow.

Module Contents#

class gensbi.recipes.flow_pipeline.ConditionalFlowPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, params=None, training_config=None, structured_obs=False, structured_cond=False)[source]#

Bases: gensbi.recipes.pipeline.AbstractPipeline

Max-likelihood NPE pipeline wrapping an MAFlow.

Parameters:

  • model (MAFlow) – A pre-built flow (e.g. MAFlow(MAFlowParams(rngs=rngs, dim=dim_obs, cond_dim=dim_cond))).

  • train_dataset (iterable) – Yield (obs, cond) batches. Shape is (B, dim, C) for each variable (C = 1 for the tabular path; see ch_obs/ch_cond).

  • val_dataset (iterable) – Yield (obs, cond) batches. Shape is (B, dim, C) for each variable (C = 1 for the tabular path; see ch_obs/ch_cond).

  • dim_obs (int)

  • dim_cond (int)

  • ch_obs (int, optional) – Channel count for the obs and cond variables. Default 1 (tabular SBI). Values > 1 enable the (B, dim, C) channel-passthrough path: the channel axis is preserved and forwarded to the flow unchanged (the flow must be built with matching channels/cond_channels in MAFlowParams).

  • ch_cond (int, optional) – Channel count for the obs and cond variables. Default 1 (tabular SBI). Values > 1 enable the (B, dim, C) channel-passthrough path: the channel axis is preserved and forwarded to the flow unchanged (the flow must be built with matching channels/cond_channels in MAFlowParams).

  • structured_obs (bool, optional) – If True, the modeled variable / condition keeps its native structured shape (the model owns it) instead of the tabular (B, dim, 1) layout. Default False.

  • structured_cond (bool, optional) – If True, the modeled variable / condition keeps its native structured shape (the model owns it) instead of the tabular (B, dim, 1) layout. Default False.

Notes

Every single-observation method (sample(), log_prob(), get_sampler(), get_log_prob_fn()) expects x_o to carry a leading batch axis (size 1 for one observation) and a channel axis: shape (1, dim_cond, C) for tabular, or (1,) + per_obs_shape for structured. A bare (B, dim) tensor is rejected — add [..., None] for C = 1. A batch axis > 1 raises ValueError — pass a batch to sample_batched() instead.

abstractmethod _make_model(params)[source]#

Create and return the model to be trained.

_prep_cond(x)[source]#
_prep_obs(x)[source]#
_wrap_model()[source]#

Wrap the model for evaluation (either using JointWrapper or ConditionalWrapper).

fit_standardization(obs_data, axis=0)[source]#

Fit the Standardize bijection buffers from training observations.

Computes per-dimension mean and standard deviation of obs_data and stores them as buffers on both the live model and the EMA model. EMA only averages Param variables, so the non-Param buffers must be set explicitly here. Must be called before train() when input standardization is desired.

Parameters:

  • obs_data (Array) – Training observations of shape (N, dim_obs) or (N, dim_obs, 1) (the autoregressive target; e.g. theta for NPE). For multichannel flows (ch_obs > 1) the shape is (N, dim_obs, C) and axis=(0, 1) yields per-channel stats.

  • axis (int or tuple of int, optional) – Reduction axis or axes for the mean/std computation. Default is 0 (per-dimension stats over the batch), which is the correct choice for the tabular (C == 1) path. Pass axis=(0, 1) for per-channel standardization when C > 1.

classmethod get_default_params(*args, **kwargs)[source]#
Abstractmethod:

Not implemented: the flow pipeline takes a pre-built MAFlow.

Raises:

NotImplementedError – Always. There are no default model params to return; construct an MAFlow directly and pass it as model=.

get_log_prob_fn(x_o, use_ema=True, **kwargs)[source]#

Return a log-probability closure for a single conditioning observation.

Parameters:

  • x_o (Array) – Single conditioning observation carrying a leading batch axis of size 1 (see get_sampler() for the shape convention). A leading batch axis > 1 raises ValueError.

  • use_ema (bool, optional) – If True (default), use the EMA model.

Returns:

log_prob_fn – A function (x_1) -> Array of shape (B,) evaluating the conditional log-probability log q(x_1 | x_o) for a batch of B parameter vectors. x_1 has shape (B, dim_obs) or (B, dim_obs, 1) on the tabular path, or (B, dim_obs, C) when ch_obs > 1 (channel-passthrough).

Return type:

Callable

get_loss_fn()[source]#

Return the max-likelihood loss function for training.

Returns a closure loss_fn(model, batch, key) -> Array that computes the mean negative log-likelihood -mean(log q(obs | cond)). batch = (obs, cond) with each element of shape (B, dim, 1). NPE convention: obs = theta, cond = x. The key argument is accepted for interface compatibility but is unused.

Returns:

loss_fn – A function (model, batch, key) -> Array returning the scalar mean negative log-likelihood.

Return type:

Callable

get_sampler(x_o, use_ema=True, **kwargs)[source]#

Return a sampler closure for a single conditioning observation.

Parameters:

  • x_o (Array) – Single conditioning observation. Must carry a leading batch axis and a channel axis for tabular cond: shape (1, dim_cond, C). For structured cond: (1,) + per_observation_shape. A leading batch axis > 1 raises ValueError (use sample_batched() for many conditions).

  • use_ema (bool, optional) – If True (default), use the EMA model; otherwise use the live model.

Returns:

sampler – A function (key, nsamples) -> Array returning the model’s native output shape (nsamples, dim_obs, C) (channel always carried).

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)[source]#
Abstractmethod:

Not implemented: the flow pipeline requires a pre-built model.

Raises:

NotImplementedError – Always. Construct an MAFlow and pass it as model= to the pipeline constructor instead.

log_prob(x_1, x_o, use_ema=True, **kwargs)[source]#

Evaluate the conditional log-probability for a batch of samples.

Parameters:

  • x_1 (Array) – Batch of parameter vectors of shape (B, dim_obs) or (B, dim_obs, 1).

  • x_o (Array) – Single conditioning observation carrying a leading batch axis of size 1 (see get_sampler() for the shape convention). A leading batch axis > 1 raises ValueError.

  • use_ema (bool, optional) – If True (default), use the EMA model.

Returns:

log_prob – Log-probabilities of shape (B,).

Return type:

Array

sample(key, x_o, nsamples=10000, use_ema=True, chunk_size=None, show_progress_bars=True, **kwargs)[source]#

Draw posterior samples for a single conditioning observation.

Parameters:

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

  • x_o (Array) – Single conditioning observation carrying a leading batch axis of size 1 (see get_sampler() for the shape convention). A leading batch axis > 1 raises ValueError.

  • nsamples (int, optional) – Number of posterior samples to draw. Default is 10 000.

  • use_ema (bool, optional) – If True (default), use the EMA model.

  • chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior. Set it to bound memory when drawing many samples from a deep flow.

  • show_progress_bars (bool, optional) – Show a progress bar over chunks (only when chunking is active). Default is True.

Returns:

samples – Posterior samples of shape (nsamples, dim_obs, 1) for the tabular default (C = 1), or (nsamples, dim_obs, C) for ch_obs = C — the channel axis is always carried for a vector-modeled variable regardless of structured_cond (a structured condition changes only x_o’s expected shape, not the modeled variable’s). When structured_obs=True, samples instead have shape (nsamples,) + per_obs_shape, the model’s native structured output.

Return type:

Array

sample_batched(key, x_o, nsamples=10000, *, use_ema=True, chunk_size=None, show_progress_bars=True, **kwargs)[source]#

Draw posterior samples for a batch of conditioning observations.

Each condition is repeated nsamples times and concatenated into a single flattened (B * nsamples, ...) batch. Without chunk_size the whole batch runs in one autoregressive pass (memory scales with B * nsamples); with chunk_size the flattened batch is sliced into pieces of at most chunk_size rows per flow.sample call — chunk boundaries may fall inside a condition, which is fine because every row is independent.

Parameters:

  • key (jax.random.PRNGKey) – Random key for the batched sampling pass.

  • x_o (Array) – Batch of observations. For tabular cond: shape (B, dim_cond, C) (a bare (B, dim_cond) raises ValueError — add a trailing channel axis). For structured cond: (B,) + per_obs_shape.

  • nsamples (int, optional) – Number of posterior samples per observation. Default is 10 000.

  • use_ema (bool, optional) – If True (default), use the EMA model.

  • chunk_size (int, optional) – Maximum number of rows of the flattened B * nsamples batch per device call. None (default) keeps the historical single-pass behavior.

  • show_progress_bars (bool, optional) – Show a progress bar over chunks (only when chunking is active). Default is True.

  • **kwargs (dict, optional) – Extra keyword arguments accepted for interface compatibility and ignored with a warning (e.g. solver arguments from AbstractPipeline).

Returns:

samples – Posterior samples of shape (nsamples, B, dim_obs, 1) for the tabular default (C = 1), or (nsamples, B, dim_obs, C) for ch_obs = C. When structured_obs=True, samples instead have shape (nsamples, B) + per_obs_shape. In both cases out[:, i] is the samples for condition i.

Return type:

Array

train(rngs, nsteps=None, save_model=True)[source]#

Train the flow model, warning if standardization was skipped.

Delegates to AbstractPipeline.train() after checking that fit_standardization() was called.

Parameters:

  • rngs (nnx.Rngs) – Random number generators for training and validation steps.

  • nsteps (int or None, optional) – Number of training steps. If None, taken from training_config["nsteps"]. Default is None.

  • save_model (bool, optional) – If True (default), serialise the model to disk after training.

Returns:

  • loss_array (list) – Per-step training losses.

  • val_loss_array (list) – Validation losses recorded at each validation checkpoint.

_standardized = False[source]#
structured_cond = False[source]#
structured_obs = False[source]#
gensbi.recipes.flow_pipeline._warn_unused_kwargs(kwargs)[source]#

Warn that solver-style kwargs are ignored by the (solver-free) flow.

The flow pipeline mirrors the flow-matching surface (which accepts **sampler_kwargs), but a normalizing flow has no ODE/SDE solver, so arguments like step_size/nsteps/solver do not apply and are silently ignored apart from this warning.