gensbi.recipes

Cookie cutter modules for creating and training SBI models.

Submodules

• gensbi.recipes.conditional_pipeline
• gensbi.recipes.flow_pipeline
• gensbi.recipes.flux1
• gensbi.recipes.flux1joint
• gensbi.recipes.id_strategies
• gensbi.recipes.joint_pipeline
• gensbi.recipes.pipeline
• gensbi.recipes.simformer
• gensbi.recipes.unconditional_pipeline
• gensbi.recipes.utils

Classes

ConditionalFlowPipeline	Max-likelihood NPE pipeline wrapping an MAFlow.
ConditionalPipeline	Model-agnostic conditional pipeline parameterized by a GenerativeMethod.
Flux1DiffusionPipeline	Model-agnostic conditional pipeline parameterized by a GenerativeMethod.
Flux1FlowPipeline	Model-agnostic conditional pipeline parameterized by a GenerativeMethod.
Flux1JointDiffusionPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
Flux1JointFlowPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
Flux1JointSMPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
Flux1SMPipeline	Model-agnostic conditional pipeline parameterized by a GenerativeMethod.
HealpixRope	Spherical RoPE ids for tokens on a HEALPix grid (name: "healpix-rope").
IdStrategy	Structural interface for pipeline id-builder strategy objects.
JointPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
SimformerDiffusionPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
SimformerFlowPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
SimformerSMPipeline	Model-agnostic joint pipeline parameterized by a GenerativeMethod.
UnconditionalPipeline	Model-agnostic unconditional pipeline parameterized by a GenerativeMethod.

Package Contents

class gensbi.recipes.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)

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)

Create and return the model to be trained.

_prep_cond(x)

_prep_obs(x)

_wrap_model()

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

fit_standardization(obs_data, axis=0)

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)

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)

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()

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)

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)

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)

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)

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)

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)

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

structured_cond = False

structured_obs = False

class gensbi.recipes.ConditionalPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, method, ch_obs=1, ch_cond=1, id_embedding_strategy=('absolute', 'absolute'), size=2, params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the ConditionalWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding (obs, cond) batches.

• val_dataset (iterable) – Validation dataset yielding (obs, cond) batches.

• dim_obs (int or tuple of int) – Dimension of the observation/parameter space.

• dim_cond (int or tuple of int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per observation token. Default is 1.

• ch_cond (int, optional) – Number of channels per conditioning token. Default is 1.

• id_embedding_strategy (tuple of (str or IdStrategy), optional) – Per-stream (obs, cond) id-builder strategy. Strings pick a built-in 1D/2D grid builder; an IdStrategy object (e.g. HealpixRope) builds ids from its own geometry. Default is ("absolute", "absolute"). NOTE this pipeline-side vocabulary is distinct from the model-side id_embedding_strategy (e.g. Flux1Params), where “rope” means “apply RoPE to the provided ids”; a HealpixRope pipeline strategy pairs with model-side ("absolute", "rope") and a 3-entry axes_dim.

• size (int or tuple of int, optional) – Patch edge length for 2D ID-embedding strategies. Default is 2. A single int is broadcast to both obs and cond (8 -> (8, 8)). A length-2 tuple (obs_size, cond_size) lets the two inputs differ. Use 1 to disable patchification for an input. Ignored for 1D strategies ("absolute", "pos1d", "rope1d").

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration. If None, uses defaults augmented by method.get_extra_training_config().

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = ConditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=5, dim_cond=3,
...     method=FlowMatchingMethod(),
... )

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

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

classmethod get_default_params(*args, **kwargs)

Abstractmethod:

get_log_prob_fn(x_o, use_ema=True, model_extras=None, **kwargs)

Get a log-probability function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• model_extras (dict, optional) – Additional model extras. Cannot override protected keys.

• **kwargs – Forwarded to method.build_log_prob_fn.

Returns:

log_prob_fn(x_1) -> log_prob

Return type:

Callable

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, use_ema=True, model_extras=None, **sampler_kwargs)

Get a sampler function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• model_extras (dict, optional) – Additional keyword arguments passed to the model during sampling (e.g. {"edge_mask": mask}). Cannot override the protected keys cond, obs_ids, cond_ids.

• **sampler_kwargs – Forwarded to method.build_sampler_fn (e.g. step_size, nsteps, solver, time_grid).

Returns:

sampler(key, nsamples) -> samples

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)

Abstractmethod:

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

log_prob(x_1, x_o, use_ema=True, *, key=None, **kwargs)

Compute log-probability of x_1 given x_o.

Parameters:

• x_1 (array-like) – Data samples to evaluate.

• x_o (array-like) – Conditioning variable.

• use_ema (bool, optional) – Use the EMA model. Default is True.

• key (jax.random.PRNGKey, optional) – Required when exact_divergence=False (Hutchinson).

• **kwargs – Forwarded to get_log_prob_fn().

Returns:

Log-probabilities.

Return type:

Array

sample(key, x_o, nsamples=10000, use_ema=True, chunk_size=None, show_progress_bars=True, **sampler_kwargs)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• 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 large model.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

loss_obj

method

path

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

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the ConditionalWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding (obs, cond) batches.

• val_dataset (iterable) – Validation dataset yielding (obs, cond) batches.

• dim_obs (int or tuple of int) – Dimension of the observation/parameter space.

• dim_cond (int or tuple of int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per observation token. Default is 1.

• ch_cond (int, optional) – Number of channels per conditioning token. Default is 1.

• id_embedding_strategy (tuple of (str or IdStrategy), optional) – Per-stream (obs, cond) id-builder strategy. Strings pick a built-in 1D/2D grid builder; an IdStrategy object (e.g. HealpixRope) builds ids from its own geometry. Default is ("absolute", "absolute"). NOTE this pipeline-side vocabulary is distinct from the model-side id_embedding_strategy (e.g. Flux1Params), where “rope” means “apply RoPE to the provided ids”; a HealpixRope pipeline strategy pairs with model-side ("absolute", "rope") and a 3-entry axes_dim.

• size (int or tuple of int, optional) – Patch edge length for 2D ID-embedding strategies. Default is 2. A single int is broadcast to both obs and cond (8 -> (8, 8)). A length-2 tuple (obs_size, cond_size) lets the two inputs differ. Use 1 to disable patchification for an input. Ignored for 1D strategies ("absolute", "pos1d", "rope1d").

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration. If None, uses defaults augmented by method.get_extra_training_config().

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = ConditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=5, dim_cond=3,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ema_model

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

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the ConditionalWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding (obs, cond) batches.

• val_dataset (iterable) – Validation dataset yielding (obs, cond) batches.

• dim_obs (int or tuple of int) – Dimension of the observation/parameter space.

• dim_cond (int or tuple of int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per observation token. Default is 1.

• ch_cond (int, optional) – Number of channels per conditioning token. Default is 1.

• id_embedding_strategy (tuple of (str or IdStrategy), optional) – Per-stream (obs, cond) id-builder strategy. Strings pick a built-in 1D/2D grid builder; an IdStrategy object (e.g. HealpixRope) builds ids from its own geometry. Default is ("absolute", "absolute"). NOTE this pipeline-side vocabulary is distinct from the model-side id_embedding_strategy (e.g. Flux1Params), where “rope” means “apply RoPE to the provided ids”; a HealpixRope pipeline strategy pairs with model-side ("absolute", "rope") and a 3-entry axes_dim.

• size (int or tuple of int, optional) – Patch edge length for 2D ID-embedding strategies. Default is 2. A single int is broadcast to both obs and cond (8 -> (8, 8)). A length-2 tuple (obs_size, cond_size) lets the two inputs differ. Use 1 to disable patchification for an input. Ignored for 1D strategies ("absolute", "pos1d", "rope1d").

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration. If None, uses defaults augmented by method.get_extra_training_config().

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = ConditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=5, dim_cond=3,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ema_model

class gensbi.recipes.Flux1JointDiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_obs = 1

dim_joint

ema_model

class gensbi.recipes.Flux1JointFlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_obs = 1

dim_joint

ema_model

class gensbi.recipes.Flux1JointSMPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, sde_type='VP', params=None, training_config=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

• sde_type (str)

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1Joint model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor (e.g. sde_type="VE" for score matching).

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ch_obs = 1

dim_joint

ema_model

class gensbi.recipes.Flux1SMPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, ch_cond=1, sde_type='VP', params=None, training_config=None)

Bases: gensbi.recipes.conditional_pipeline.ConditionalPipeline

Model-agnostic conditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the ConditionalWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding (obs, cond) batches.

• val_dataset (iterable) – Validation dataset yielding (obs, cond) batches.

• dim_obs (int or tuple of int) – Dimension of the observation/parameter space.

• dim_cond (int or tuple of int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per observation token. Default is 1.

• ch_cond (int, optional) – Number of channels per conditioning token. Default is 1.

• id_embedding_strategy (tuple of (str or IdStrategy), optional) – Per-stream (obs, cond) id-builder strategy. Strings pick a built-in 1D/2D grid builder; an IdStrategy object (e.g. HealpixRope) builds ids from its own geometry. Default is ("absolute", "absolute"). NOTE this pipeline-side vocabulary is distinct from the model-side id_embedding_strategy (e.g. Flux1Params), where “rope” means “apply RoPE to the provided ids”; a HealpixRope pipeline strategy pairs with model-side ("absolute", "rope") and a 3-entry axes_dim.

• size (int or tuple of int, optional) – Patch edge length for 2D ID-embedding strategies. Default is 2. A single int is broadcast to both obs and cond (8 -> (8, 8)). A length-2 tuple (obs_size, cond_size) lets the two inputs differ. Use 1 to disable patchification for an input. Ignored for 1D strategies ("absolute", "pos1d", "rope1d").

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration. If None, uses defaults augmented by method.get_extra_training_config().

• sde_type (str)

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = ConditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=5, dim_cond=3,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Flux1 model to be trained.

classmethod get_default_params(dim_obs, dim_cond, ch_obs, ch_cond)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor (e.g. sde_type="VE" for score matching).

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

ema_model

class gensbi.recipes.HealpixRope

Spherical RoPE ids for tokens on a HEALPix grid (name: “healpix-rope”).

Wraps gensbi.recipes.utils.init_ids_healpix() (see there for the method and its rationale) with the geometry needed to build the ids — which the string-enum API cannot carry, since the pipeline only passes a token count.

Parameters:

• nside (int) – HEALPix resolution of the token grid (power of 2).

• base_pixels (sequence of int, optional) – Base pixels (0..11) covered by the grid; None = full sky.

__post_init__()

build(dim)

Return (ids, num_tokens); dim must match the grid.

base_pixels: Tuple[int, Ellipsis] | Sequence[int] | None = None

name: ClassVar[str] = 'healpix-rope'

nside: int

property num_tokens: int

n_faces * nside**2.

Type:

Token count of the grid

Return type:

int

property theta: int

healpix_rope_theta().

Type:

Suggested model-side RoPE theta

Return type:

int

class gensbi.recipes.IdStrategy

Bases: Protocol

Structural interface for pipeline id-builder strategy objects.

Any object with a name and a build(dim) -> (ids, resolved_dim) method can be passed in an id_embedding_strategy tuple slot; the pipeline calls build with the corresponding dim_obs/dim_cond. Strategies own their full geometry — unlike the string strategies they receive no semantic_id/size.

build(dim)

Return (ids, resolved_dim) for a stream of dim tokens.

name: str

class gensbi.recipes.JointPipeline(model, train_dataset, val_dataset, dim_obs, dim_cond, method, ch_obs=1, condition_mask_kind='structured', params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

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

classmethod get_default_params(*args, **kwargs)

Abstractmethod:

get_log_prob_fn(x_o, use_ema=True, prior=None, model_extras=None, **kwargs)

Get a log-probability function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• prior (numpyro.distributions.Distribution, optional) –

  Obs-space prior for log-probability evaluation. The method’s prior lives on the full joint space (dim_joint, ch) and cannot be automatically marginalized for arbitrary priors.

  • Default Gaussian: auto-constructed — no need to provide.

  • Custom prior: must supply the correct obs-space marginal.

• model_extras (dict, optional) – Additional model extras. Cannot override protected keys.

• **kwargs – Forwarded to method.build_log_prob_fn.

Returns:

log_prob_fn(x_1) -> log_prob

Return type:

Callable

Raises:

ValueError – If the joint prior is non-Gaussian and no prior is provided.

get_loss_fn()

Return the loss function for training/validation.

get_sampler(x_o, use_ema=True, model_extras=None, **sampler_kwargs)

Get a sampler function.

Parameters:

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

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• model_extras (dict, optional) – Additional keyword arguments passed to the model during sampling (e.g. {"edge_mask": mask}). Cannot override the protected keys cond, obs_ids, cond_ids.

• **sampler_kwargs – Forwarded to method.build_sampler_fn.

Returns:

sampler(key, nsamples) -> samples

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)

Abstractmethod:

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

log_prob(x_1, x_o, use_ema=True, prior=None, *, key=None, **kwargs)

Compute log-probability of x_1 given x_o.

Parameters:

• x_1 (array-like) – Data samples to evaluate.

• x_o (array-like) – Conditioning variable.

• use_ema (bool, optional) – Use the EMA model. Default is True.

• prior (numpyro.distributions.Distribution, optional) – Obs-space prior distribution. See get_log_prob_fn() for details.

• key (jax.random.PRNGKey, optional) – Required when exact_divergence=False (Hutchinson).

• **kwargs – Forwarded to get_log_prob_fn().

Returns:

Log-probabilities.

Return type:

Array

sample(key, x_o, nsamples=10000, use_ema=True, chunk_size=None, show_progress_bars=True, **sampler_kwargs)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

condition_mask_kind = 'structured'

dim_joint

loss_obj

method

path

class gensbi.recipes.SimformerDiffusionPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, edge_mask=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Simformer model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

sample(key, x_o, nsamples=10000, nsteps=18, use_ema=True, return_intermediates=False, chunk_size=None, show_progress_bars=True)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

ch_obs = 1

dim_joint

edge_mask = None

ema_model

class gensbi.recipes.SimformerFlowPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, params=None, training_config=None, edge_mask=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Simformer model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor.

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

sample(key, x_o, nsamples=10000, step_size=0.01, use_ema=True, time_grid=None, chunk_size=None, show_progress_bars=True)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

ch_obs = 1

dim_joint

edge_mask = None

ema_model

class gensbi.recipes.SimformerSMPipeline(train_dataset, val_dataset, dim_obs, dim_cond, ch_obs=1, sde_type='VP', params=None, training_config=None, edge_mask=None, condition_mask_kind='structured')

Bases: gensbi.recipes.joint_pipeline.JointPipeline

Model-agnostic joint pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the JointWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding concatenated x_1 batches (obs and cond concatenated along the token dimension).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the observation/parameter space.

• dim_cond (int) – Dimension of the conditioning space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• condition_mask_kind (str, optional) – Kind of condition mask. One of "structured" or "posterior". Default is "structured".

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

• sde_type (str)

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = JointPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=2, dim_cond=7,
...     method=FlowMatchingMethod(),
... )

_make_model(params)

Create and return the Simformer model to be trained.

classmethod get_default_params(dim_joint, in_channels)

Return a dictionary of default model parameters.

classmethod init_pipeline_from_config(train_dataset, val_dataset, dim_obs, dim_cond, config_path, checkpoint_dir, **kwargs)

Initialize the pipeline from a configuration file.

Parameters:

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

• **kwargs – Additional keyword arguments forwarded to the constructor (e.g. sde_type="VE" for score matching).

• dim_obs (int)

• dim_cond (int)

• checkpoint_dir (str)

sample(key, x_o, nsamples=10000, nsteps=1000, use_ema=True, return_intermediates=False, chunk_size=None, show_progress_bars=True)

Draw samples from the model.

Parameters:

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

• x_o (array-like) – Conditioning variable.

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

ch_obs = 1

dim_joint

edge_mask = None

ema_model

class gensbi.recipes.UnconditionalPipeline(model, train_dataset, val_dataset, dim_obs, method, ch_obs=1, params=None, training_config=None)

Bases: gensbi.recipes.pipeline.AbstractPipeline

Model-agnostic unconditional pipeline parameterized by a GenerativeMethod.

Unlike the old method-specific pipeline classes, this class works with any generative method and any user-provided model that conforms to the UnconditionalWrapper interface.

Parameters:

• model (nnx.Module) – The model to be trained.

• train_dataset (iterable) – Training dataset yielding x_1 batches (not tuples).

• val_dataset (iterable) – Validation dataset.

• dim_obs (int) – Dimension of the data space.

• method (GenerativeMethod) – Strategy object (e.g. FlowMatchingMethod(), DiffusionEDMMethod(), ScoreMatchingMethod()).

• ch_obs (int, optional) – Number of channels per token. Default is 1.

• params (optional) – Model parameters (stored but not used directly).

• training_config (dict, optional) – Training configuration.

Examples

>>> from gensbi.core import FlowMatchingMethod
>>> pipeline = UnconditionalPipeline(
...     model=my_model,
...     train_dataset=train_ds,
...     val_dataset=val_ds,
...     dim_obs=9,
...     method=FlowMatchingMethod(),
... )

abstractmethod _make_model()

Create and return the model to be trained.

_wrap_model()

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

classmethod get_default_params(*args, **kwargs)

Abstractmethod:

get_log_prob_fn(use_ema=True, **kwargs)

Get a log-probability function.

Parameters:

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• **kwargs – Forwarded to method.build_log_prob_fn.

Returns:

log_prob_fn(x_1) -> log_prob

Return type:

Callable

get_loss_fn()

Return the loss function for training/validation.

get_sampler(use_ema=True, **sampler_kwargs)

Get a sampler function.

Parameters:

• use_ema (bool, optional) – Whether to use the EMA model. Default is True.

• **sampler_kwargs – Forwarded to method.build_sampler_fn.

Returns:

sampler(key, nsamples) -> samples

Return type:

Callable

classmethod init_pipeline_from_config(*args, **kwargs)

Abstractmethod:

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

log_prob(x_1, use_ema=True, *, key=None, **kwargs)

Compute log-probability of x_1.

Parameters:

• x_1 (array-like) – Data samples to evaluate.

• use_ema (bool, optional) – Use the EMA model. Default is True.

• key (jax.random.PRNGKey, optional) – Required when exact_divergence=False (Hutchinson).

• **kwargs – Forwarded to get_log_prob_fn().

Returns:

Log-probabilities.

Return type:

Array

sample(key, nsamples=10000, use_ema=True, chunk_size=None, show_progress_bars=True, **sampler_kwargs)

Draw samples from the model.

Parameters:

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

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

• use_ema (bool, optional) – Use the EMA model. Default is True.

• chunk_size (int, optional) – Maximum number of samples drawn per device call. None (default) draws everything in one call — identical to the historical behavior.

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

• **sampler_kwargs – Forwarded to get_sampler().

Returns:

Samples of shape (nsamples, dim_obs, ch_obs).

Return type:

Array

abstractmethod sample_batched(*args, **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

loss_obj

method

path