gensbi.recipes.pipeline

Pipeline module for GenSBI.

This module provides an abstract pipeline class for training and evaluating conditional generative models (such as conditional flow matching or diffusion models) in the GenSBI framework. It handles model creation, training loop, optimizer setup, checkpointing, and evaluation utilities.

For practical implementations, subclasses should implement specific model architectures, loss functions, and sampling methods. See JointPipeline and ConditionalPipeline for concrete examples.

Classes

AbstractPipeline	Abstract base class for GenSBI training pipelines.
ModelEMA	Exponential Moving Average (EMA) optimizer for maintaining a smoothed version of model parameters.

Functions

_cast_state_to_target_dtypes(state, target_state)	Cast state's leaves to match target_state's dtypes, in place.
_chunked_draw(sampler, key, nsamples, chunk_size[, ...])	Draw nsamples from sampler in memory-bounded chunks.
_sample_concat_axis(sampler_kwargs)	Axis carrying the sample dimension in a sampler's output.
_validate_chunk_size(chunk_size)	Reject non-positive chunk sizes before any chunk arithmetic.
_warn_if_not_fp32_master_weights(model)	Warn when trainable params are not fp32 master weights.
ema_step(ema_model, model, ema_optimizer)	Update EMA model with current model parameters.

Module Contents

class gensbi.recipes.pipeline.AbstractPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=None, params=None, training_config=None)

Bases: abc.ABC

Abstract base class for GenSBI training pipelines.

This class provides a template for implementing training and evaluation pipelines for conditional generative models. Subclasses should implement model creation, default parameter setup, loss function, sampling, and evaluation methods.

Parameters:

• train_dataset (iterable) – Training dataset, should yield batches of data.

• val_dataset (iterable) – Validation dataset, should yield batches of data.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• model (nnx.Module, optional) – The model to be trained. If None, the model is created using _make_model.

• params (dict, optional) – Model parameters. If None, uses defaults from _get_default_params.

• ch_obs (int, optional) – Number of channels in the observation data. Default is 1.

• ch_cond (int, optional) – Number of channels in the conditional data (if applicable). Default is None.

• training_config (dict, optional) – Training configuration. If None, uses defaults from get_default_training_config.

_get_ema_optimizer()

Construct the EMA optimizer for maintaining an exponential moving average of model parameters. :returns: ema_optimizer – The EMA optimizer instance. :rtype: ModelEMA

_get_optimizer()

Construct the optimizer for training, including learning rate scheduling and gradient clipping.

Returns:

optimizer – The optimizer instance for the model.

Return type:

nnx.Optimizer

abstractmethod _make_model(params)

Create and return the model to be trained.

_restore_best_state(best_state, best_state_ema)

Restore the best model and EMA states (used after early stopping).

Parameters:

• best_state (nnx.State) – Best model state recorded during training.

• best_state_ema (nnx.State) – Best EMA model state recorded during training.

_run_validation(val_step, batch_val, rng_val, min_val, best_state, best_state_ema, counter, val_error_ratio, loss_array, val_loss_array, l_train)

Run a validation step and update early-stopping bookkeeping.

Parameters:

• val_step (Callable) – Validation step function.

• batch_val (Any) – Fixed validation batch.

• rng_val (jax.random.PRNGKey) – Fixed validation RNG key.

• min_val (float) – Best validation loss seen so far.

• best_state (nnx.State) – Current best model state.

• best_state_ema (nnx.State) – Current best EMA model state.

• counter (int) – Early-stopping patience counter.

• val_error_ratio (float) – Threshold ratio for incrementing the counter.

• loss_array (list) – Training loss history (mutated in place).

• val_loss_array (list) – Validation loss history (mutated in place).

• l_train (float) – Current smoothed training loss.

Returns:

• l_val (float) – Validation loss for this step.

• ratio (float) – Ratio of current validation loss to best.

• min_val (float) – Updated best validation loss.

• best_state (nnx.State) – Updated best model state.

• best_state_ema (nnx.State) – Updated best EMA model state.

• counter (int) – Updated early-stopping counter.

abstractmethod _wrap_model()

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

export_safetensors(path, *, ema=True, metadata=None)

Export trained weights to a single .safetensors file.

ema=True (default) exports the EMA model – usually the weights you want for inference and for sharing. Pass ema=False for the primary model. This is a thin wrapper over gensbi.utils.serialization.save_safetensors().

abstractmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)

classmethod get_default_training_config()

Return a dictionary of default training configuration parameters.

Returns:

training_config – Default training configuration.

Return type:

dict

abstractmethod get_log_prob_fn(*args, **kwargs)

Get a log-probability function for evaluating data under the model.

Returns:

log_prob_fn – (x_1) -> log_prob

Return type:

Callable

abstractmethod get_loss_fn()

Return the loss function for training/validation.

abstractmethod get_sampler(key, x_o, step_size=0.01, use_ema=True, time_grid=None, **model_extras)

Get a sampler function for generating samples from the trained model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

• step_size (float, optional) – Step size for the sampler.

• use_ema (bool, optional) – Whether to use the EMA model for sampling.

• time_grid (array-like, optional) – Time grid for the sampler (if applicable).

• model_extras (dict, optional) – Additional model-specific parameters.

Returns:

sampler – A function that generates samples when called with a random key and number of samples.

Return type:

Callable: key, nsamples -> samples

get_train_step_fn(loss_fn)

Return the training step function, which performs a single optimization step.

Returns:

train_step – JIT-compiled training step function.

Return type:

Callable

get_val_step_fn(loss_fn)

Return the validation step function, which performs a single optimization step.

Returns:

val_step – JIT-compiled validation step function.

Return type:

Callable

import_safetensors(path, *, ema=True, strict=True)

Load weights from a .safetensors file into this pipeline in place.

ema=True (default) loads into the EMA model. Thin wrapper over gensbi.utils.serialization.load_safetensors().

abstractmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir)

Initialize the pipeline from a configuration file.

Parameters:

• train_dataset (iterable) – Training dataset.

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimensionality of the parameter (theta) space.

• dim_cond (int) – Dimensionality of the observation (x) space.

• config_path (str) – Path to the configuration file.

• checkpoint_dir (str) – Directory for saving checkpoints.

Returns:

pipeline – An instance of the pipeline initialized from the configuration.

Return type:

AbstractPipeline

abstractmethod log_prob(x_1, *args, **kwargs)

Compute log-probability of data under the trained model.

Returns:

Log-probabilities.

Return type:

Array

restore_model(experiment_id=None)

Restore model and EMA model from checkpoints.

Parameters:

experiment_id (str, optional) – Experiment identifier. If None, uses training_config value.

abstractmethod sample(key, x_o, nsamples=10000)

Generate samples from the trained model.

Parameters:

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

• x_o (array-like) – Conditioning variable (e.g., observed data).

• nsamples (int, optional) – Number of samples to generate.

Returns:

samples – Generated samples of size (nsamples, dim_obs, ch_obs).

Return type:

array-like

sample_batched(key, x_o, nsamples, *args, chunk_size=None, show_progress_bars=True, **kwargs)

Generate samples from the trained model in batches.

Loops over the B conditions in x_o one at a time and, when chunk_size is set, additionally draws each condition’s samples in memory-bounded chunks of at most chunk_size samples per device call.

Parameters:

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

• x_o (array-like) – Conditioning variable (e.g., observed data), leading batch axis of size B.

• nsamples (int) – Number of samples to generate per condition.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws all nsamples for a condition in a single call — identical to the historical behavior.

• show_progress_bars (bool, optional) – Whether to display a progress bar over the B * n_chunks device calls. Default is True.

• args (tuple) – Additional positional arguments for the sampler.

• kwargs (dict) – Additional keyword arguments for the sampler.

Returns:

samples – Generated samples of shape (nsamples, batch_size_cond, dim_obs, ch_obs).

Return type:

array-like

save_model(experiment_id=None)

Save model and EMA model checkpoints.

Parameters:

experiment_id (str, optional) – Experiment identifier. If None, uses training_config value.

train(rngs, nsteps=None, save_model=True)

Run the training loop for the model.

Parameters:

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

• nsteps (Optional[int])

Returns:

• loss_array (list) – List of training losses.

• val_loss_array (list) – List of validation losses.

Return type:

Tuple[list, list]

update_training_config(new_config)

Update the training configuration with new parameters.

Parameters:

new_config (dict) – New training configuration parameters.

ch_cond = None

ch_obs = 1

dim_cond

dim_obs

ema_model_wrapped = None

loss_fn = None

model

model_wrapped = None

p0_dist_model = None

params = None

path = None

train_dataset

train_dataset_iter

training_config = None

val_dataset

val_dataset_iter

class gensbi.recipes.pipeline.ModelEMA(model, tx)

Bases: flax.nnx.Optimizer

Exponential Moving Average (EMA) optimizer for maintaining a smoothed version of model parameters.

This optimizer keeps an exponential moving average of the model parameters, which can help stabilize training and improve evaluation performance. The EMA parameters are updated at each training step.

Parameters:

• model (nnx.Module) – The model whose parameters will be tracked.

• tx (optax.GradientTransformation) – The Optax transformation defining the EMA update rule.

update(model, model_orginal)

Update the EMA parameters using the current model parameters. :param model: The model with EMA parameters to be updated. :type model: nnx.Module :param model_orginal: The original model with current parameters. :type model_orginal: nnx.Module

Parameters:

model_orginal (flax.nnx.Module)

gensbi.recipes.pipeline._cast_state_to_target_dtypes(state, target_state)

Cast state’s leaves to match target_state’s dtypes, in place.

Old checkpoints (e.g. bf16 master weights from a pre-mixed-precision run) must still be restorable into a model whose current dtype is fp32; mirrors the arr.astype(want.dtype) loop already used by gensbi.utils.serialization.load_safetensors().

Parameters:

• state (flax.nnx.State)

• target_state (flax.nnx.State)

Return type:

flax.nnx.State

gensbi.recipes.pipeline._chunked_draw(sampler, key, nsamples, chunk_size, show_progress_bars=True, concat_axis=0, sampler_kwargs=None, pbar=None)

Draw nsamples from sampler in memory-bounded chunks.

Parameters:

• sampler (Callable) – sampler(key, nsamples, **sampler_kwargs) -> Array.

• key (jax.random.PRNGKey) – Random key. With no chunking it is passed through UNCHANGED so the result is bit-identical to calling sampler directly.

• nsamples (int) – Total number of samples to draw.

• chunk_size (int or None) – Maximum samples per sampler call. None (or any value >= nsamples) disables chunking.

• show_progress_bars (bool, optional) – Show a tqdm bar over chunks (only when chunking is active and no external pbar is supplied).

• concat_axis (int, optional) – Axis to concatenate chunks along — 0 for plain samples, 1 when the sampler returns intermediates with a leading time axis (see _sample_concat_axis()).

• sampler_kwargs (dict, optional) – Extra keyword arguments forwarded to every sampler call (e.g. {"model_extras": ...}).

• pbar (tqdm-like, optional) – External progress bar; when given it is updated once per chunk and no internal bar is created (used by sample_batched for a single bar across conditions).

Returns:

nsamples samples, concatenated along concat_axis.

Return type:

Array

gensbi.recipes.pipeline._sample_concat_axis(sampler_kwargs)

Axis carrying the sample dimension in a sampler’s output.

Solvers stack intermediates along a leading, statically-sized time axis, so chunked outputs must concatenate along axis 1 instead of 0. Intermediates are requested either explicitly (return_intermediates=True — EDM and score-matching methods) or implicitly by passing a non-None time_grid (FlowMatchingMethod.build_sampler_fn turns intermediates on for any explicit time grid).

Parameters:

sampler_kwargs (dict)

Return type:

int

gensbi.recipes.pipeline._validate_chunk_size(chunk_size)

Reject non-positive chunk sizes before any chunk arithmetic.

Parameters:

chunk_size (Optional[int])

Return type:

None

gensbi.recipes.pipeline._warn_if_not_fp32_master_weights(model)

Warn when trainable params are not fp32 master weights.

Mixed precision in GenSBI stores master weights in fp32 and selects the compute dtype via each model’s dtype knob; bf16 master weights break AdamW moment accumulation and make optax.ema unable to integrate (1 - decay)-scale updates.

gensbi.recipes.pipeline.ema_step(ema_model, model, ema_optimizer)

Update EMA model with current model parameters.

Parameters:

ema_optimizer (flax.nnx.Optimizer)