gensbi.models#

Model architectures for GenSBI.

This package provides transformer-based models for simulation-based inference, including Flux1, Simformer, and autoencoder architectures, along with their associated loss functions and wrappers.

Submodules#

Classes#

`ConditionalWrapper`	Wrapper for conditional models to handle input expansion and calling convention.
`Flux1`	Transformer model for flow matching on sequences.
`Flux1Joint`	Flux1Joint model for joint density estimation.
`Flux1JointParams`	Parameters for the Flux1Joint model.
`Flux1Params`	Parameters for the Flux1 model.
`JointWrapper`	Wrapper for joint models to handle both conditioned and unconditioned inference.
`MAFlow`	Masked Autoregressive Flow for exact density evaluation and sampling.
`MAFlowParams`	Architecture parameters for `MAFlow`.
`Simformer`	Simformer model for joint density estimation.
`SimformerParams`	Parameters for the Simformer model.
`TarFlow`	Transformer autoregressive normalizing flow density model.
`TarFlowParams`	Architecture parameters for `TarFlow`.
`UnconditionalWrapper`	Wrapper for unconditional models to handle input expansion and calling convention.

Package Contents#

class gensbi.models.ConditionalWrapper(model)[source]#

Bases: gensbi.utils.model_wrapping.ModelWrapper

Wrapper for conditional models to handle input expansion and calling convention.

Parameters:: model (The conditional model instance to wrap.)

__call__(t, obs, obs_ids, cond, cond_ids, conditioned=True, guidance=None, **kwargs)[source]#

Call the wrapped model with expanded inputs.

Parameters:

t (Array) – Time steps.
obs (Array) – Observations.
obs_ids (Array) – Observation identifiers.
cond (Array) – Conditioning values.
cond_ids (Array) – Conditioning identifiers.
conditioned (bool | Array, optional) – Whether to use conditioning. Defaults to True.
guidance (Array | None, optional) – Optional guidance input.

Returns:

Model output.

Return type:

Array

class gensbi.models.Flux1(params)[source]#

Bases: flax.nnx.Module

Transformer model for flow matching on sequences.

Parameters:: params (Flux1Params)

__call__(t, obs, obs_ids, cond, cond_ids, conditioned=True, guidance=None)[source]#

Parameters:

t (jax.Array)
obs (jax.Array)
obs_ids (jax.Array)
cond (jax.Array)
cond_ids (jax.Array)
conditioned (bool | jax.Array)
guidance (jax.Array | None)

Return type:

jax.Array

cond_in#

double_blocks#

final_layer#

hidden_size#

id_merge_mode#

in_channels#

num_heads#

obs_in#

out_channels#

params#

qkv_features#

single_blocks#

time_in#

vector_in#

class gensbi.models.Flux1Joint(params)[source]#

Bases: flax.nnx.Module

Flux1Joint model for joint density estimation.

Parameters:: params (Flux1JointParams) – Parameters for the Flux1Joint model.

__call__(t, obs, node_ids, condition_mask, guidance=None, edge_mask=None)[source]#

Parameters:

t (jax.Array)
obs (jax.Array)
node_ids (jax.Array)
condition_mask (jax.Array)
guidance (jax.Array | None)
edge_mask (Optional[jax.Array])

Return type:

jax.Array

final_layer#

hidden_size#

in_channels#

num_heads#

obs_in#

out_channels#

params#

qkv_features#

single_blocks#

time_in#

vector_in#

class gensbi.models.Flux1JointParams[source]#

Parameters for the Flux1Joint model.

GenSBI uses the tensor convention (batch, dim, channels).

For joint density estimation, the model consumes a single sequence obs that mixes all variables you want to model jointly. In this case:

dim_joint is the number of tokens in that joint sequence.
in_channels is the number of channels/features per token.

In many SBI-style problems you will still use in_channels = 1 (one scalar per token), but for some datasets a token may carry multiple features.

Parameters:

in_channels (int) – Number of channels/features per token.
vec_in_dim (Union[int, None]) – Dimension of the vector input, if applicable.
mlp_ratio (float) – Ratio for the MLP layers.
num_heads (int) – Number of attention heads.
depth_single_blocks (int) – Number of single stream blocks.
val_emb_dim (int) – Number of features per head used to embed the data.
cond_emb_dim (int) – Number of features per head used to encode the condition mask, which determines the features on which we are conditioning.
id_emb_dim (int) – Number of features per head used to encode the token ids.
qkv_bias (bool) – Whether to use bias in QKV layers.
rngs (nnx.Rngs) – Random number generators for initialization.
dim_joint (int) – Number of tokens in the joint sequence.
id_merge_mode (str) – Strategy for combining embeddings (“sum” or “concat”).
id_embedding_strategy (str) – Kind of ID embedding. Currently only “absolute” is supported for Flux1Joint.
guidance_embed (bool) – Whether to use guidance embedding.
param_dtype (DTypeLike) – Data type for master-weight storage (parameters). Defaults to jnp.float32.
dtype (DTypeLike) – Compute/matmul dtype used by the compute layers. Defaults to jnp.bfloat16.

__post_init__()[source]#

cond_emb_dim: int#

depth_single_blocks: int#

dim_joint: int#

dtype: jax.typing.DTypeLike#

guidance_embed: bool = False#

id_emb_dim: int#

id_embedding_strategy: str = 'absolute'#

id_merge_mode: str = 'sum'#

in_channels: int#

mlp_ratio: float#

num_heads: int#

param_dtype: jax.typing.DTypeLike#

qkv_bias: bool#

rngs: flax.nnx.Rngs#

val_emb_dim: int#

vec_in_dim: int | None#

class gensbi.models.Flux1Params[source]#

Parameters for the Flux1 model.

GenSBI uses the tensor convention (batch, dim, channels).

dim_* counts tokens (how many distinct observables/variables you have).
channels counts features per token (how many values each observable carries).

For conditional SBI with Flux1:

Parameters to infer (often denoted $ heta$) have shape (batch, dim_obs, in_channels).
In most SBI problems in_channels = 1 (one scalar per parameter token).
Conditioning data (often denoted $x$) has shape (batch, dim_cond, context_in_dim).
context_in_dim can be > 1 (e.g., multiple detectors or multiple features per measured token).

Data Stucture and ID Embeddings:

Flux1 supports unstructured, 1D, and 2D data (and can be extended to ND) through different ID embedding strategies. The model needs to know what each token represents distinct from its value. This is handled by id_embedding_strategy.

absolute: Learned embeddings. Use for unstructured data (order doesn’t matter, e.g. physical parameters).
Initialize IDs using gensbi.recipes.utils.init_ids_1d (the semantic_id will be ignored).
pos1d / rope1d: 1D positional embeddings. Use for sequential data (order matters, e.g. time series, spectra).
Initialize IDs using gensbi.recipes.utils.init_ids_1d. The semantic_id is optional for pos1d but recommended for rope1d.
pos2d / rope2d: 2D positional embeddings. Use for image data or 2D grids.
Initialize IDs using gensbi.recipes.utils.init_ids_2d. The semantic_id is optional for pos2d but recommended for rope2d.

Combining ID Embeddings:

Strategies for combining the value and ID embeddings (id_merge_mode):

“sum” (Default): The value and ID embeddings are summed. This is the standard approach for large transformers. Requires axes_dim to be specified. Recommended for: Large models, high-dimensional data, or when using RoPE.
“concat”: The value and ID embeddings are concatenated. Requires val_emb_dim (features for value) and id_emb_dim (features for ID) to be specified. Recommended for: Small models (low dimension per head, few heads) to reduce confusion between value and positional information. A good starting ratio for val_emb_dim : id_emb_dim is 1:1.

Preprocessing for Images/2D Data:

Patchification: 2D images must be patchified (flattened into a sequence of tokens) before passing them to the model. Use gensbi.models.core.patching.patchify_2d for this purpose.
Normalization: To speed up convergence, ensure data is normalized to 0 mean and unit variance.

Note

See the documentation and tutorials for more information on id embeddings and data preprocessing.

Parameters:

in_channels (int) – Number of channels per observation/parameter token.
vec_in_dim (Union[int, None]) – Dimension of the vector input, if applicable.
context_in_dim (int) – Number of channels per conditioning token.
mlp_ratio (float) – Ratio for the MLP layers.
num_heads (int) – Number of attention heads.
depth (int) – Number of double stream blocks.
depth_single_blocks (int) – Number of single stream blocks.
qkv_bias (bool) – Whether to use bias in QKV layers.
rngs (nnx.Rngs) – Random number generators for initialization.
dim_obs (int) – Number of observation/parameter tokens.
dim_cond (int) – Number of conditioning tokens.
axes_dim (Optional[list[int]]) – Dimensions of the axes for positional encoding (required for “sum” strategy).
val_emb_dim (Optional[int]) – Features per head for value embedding (required for “concat” strategy).
id_emb_dim (Optional[int]) – Features per head for ID embedding (required for “concat” strategy).
id_merge_mode (str) – Strategy for combining embeddings (“sum” or “concat”). Default is “sum”.
theta (Optional[int]) – Scaling factor for positional encoding.
id_embedding_strategy (tuple[str, str]) – Kind of ID embedding for obs and cond respectively. Options are “absolute”, “pos1d”, “pos2d”, “rope1d”, “rope2d”.
guidance_embed (bool) – Whether to use guidance embedding.
param_dtype (DTypeLike) – Data type for master-weight storage (parameters). Defaults to jnp.float32.
dtype (DTypeLike) – Compute/matmul dtype used by the compute layers. Defaults to jnp.bfloat16.

__post_init__()[source]#

axes_dim: list[int] | None = None#

context_in_dim: int#

depth: int#

depth_single_blocks: int#

dim_cond: int#

dim_obs: int#

dtype: jax.typing.DTypeLike#

guidance_embed: bool = False#

id_emb_dim: int | None = None#

id_embedding_strategy: tuple[str, str] = ('absolute', 'absolute')#

id_merge_mode: str = 'sum'#

in_channels: int#

mlp_ratio: float#

num_heads: int#

param_dtype: jax.typing.DTypeLike#

qkv_bias: bool#

rngs: flax.nnx.Rngs#

theta: int | None = None#

val_emb_dim: int | None = None#

vec_in_dim: int | None#

class gensbi.models.JointWrapper(model)[source]#

Bases: gensbi.utils.model_wrapping.ModelWrapper

Wrapper for joint models to handle both conditioned and unconditioned inference.

Parameters:

model (The joint model instance to wrap.)
conditioned (bool, optional) – Whether to use conditioning by default. Defaults to True.

__call__(t, obs, obs_ids, cond, cond_ids, conditioned=True, **kwargs)[source]#

Call the wrapped model for either conditioned or unconditioned inference.

Parameters:

t (Array) – Time steps.
obs (Array) – Observations.
obs_ids (Array) – Observation identifiers.
cond (Array) – Conditioning values.
cond_ids (Array) – Conditioning identifiers.
conditioned (bool, optional) – Whether to use conditioning. If None, uses the default set at initialization.
**kwargs (Additional keyword arguments passed to the model.)

Returns:

Model output.

Return type:

Array

conditioned(obs, obs_ids, cond, cond_ids, t, **kwargs)[source]#

Perform conditioned inference.

Parameters:

obs (Array) – Observations.
obs_ids (Array) – Observation identifiers.
cond (Array) – Conditioning values.
cond_ids (Array) – Conditioning identifiers.
t (Array) – Time steps.
**kwargs (Additional keyword arguments passed to the model.)

Returns:

Conditioned output (only for unconditioned variables).

Return type:

Array

unconditioned(obs, obs_ids, t, **kwargs)[source]#

Perform unconditioned inference.

Parameters:

obs (Array) – Observations.
obs_ids (Array) – Observation identifiers.
t (Array) – Time steps.
**kwargs (Additional keyword arguments passed to the model.)

Returns:

Unconditioned output.

Return type:

Array

class gensbi.models.MAFlow(params)[source]#

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()[source]#

log_prob(x, cond=None)[source]#

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)[source]#

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)[source]#

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.MAFlowParams[source]#

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__()[source]#

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#

class gensbi.models.Simformer(params, embedding_net_value=None)[source]#

Bases: flax.nnx.Module

Simformer model for joint density estimation.

Parameters:

params (SimformerParams) – Parameters for the Simformer model.
embedding_net_value (Optional[flax.nnx.Module])

__call__(t, obs, node_ids, condition_mask, edge_mask=None)[source]#

Forward pass of the Simformer model.

Parameters:

t (Array) – Time steps.
obs (Array) – Input data.
args (Optional[dict]) – Additional arguments.
node_ids (Array) – Node identifiers.
condition_mask (Array) – Mask for conditioning.
edge_mask (Optional[Array]) – Mask for edges.

Returns:

Model output.

Return type:

Array

cond_emb_dim#

condition_embedding#

embedding_net_id#

embedding_time#

id_emb_dim#

in_channels#

output_fn#

params#

total_tokens#

transformer#

val_emb_dim#

class gensbi.models.SimformerParams[source]#

Parameters for the Simformer model.

GenSBI uses the tensor convention (batch, dim, channels).

For Simformer (joint modeling), the input obs is a single sequence with:

dim_joint: number of tokens in the sequence (how many variables / measured points).
in_channels: number of channels/features per token.

Conditioning is controlled via condition_mask at call time (mask is over tokens, not channels): tokens with mask=1 are treated as conditioned.

Parameters:

rngs (nnx.Rngs) – Random number generators for initialization.
in_channels (int) – Number of channels/features per token.
val_emb_dim (int) – Dimension of the value embeddings.
id_emb_dim (int) – Dimension of the ID embeddings.
cond_emb_dim (int) – Dimension of the condition embeddings.
dim_joint (int) – Number of tokens in the joint sequence.
fourier_features (int) – Number of Fourier features for time embedding.
num_heads (int) – Number of attention heads.
depth (int) – Number of transformer layers.
mlp_ratio (int) – Widening factor for the transformer parameters (MLP ratio).
qkv_features (int) – Number of features for QKV layers.
num_hidden_layers (int) – Number of hidden layers in the transformer.
param_dtype (DTypeLike) – Data type for master-weight storage (parameters). Defaults to jnp.float32.
dtype (DTypeLike) – Compute/matmul dtype used by the MLP compute layers. The attention blocks always compute in fp32 regardless of this setting (see transformer.AttentionBlock). Defaults to jnp.bfloat16.

__post_init__()[source]#

cond_emb_dim: int#

depth: int#

dim_joint: int#

dtype: jax.typing.DTypeLike#

fourier_features: int = 128#

id_emb_dim: int#

in_channels: int#

mlp_ratio: int = 3#

num_heads: int#

num_hidden_layers: int = 1#

param_dtype: jax.typing.DTypeLike#

qkv_features: int | None = None#

rngs: flax.nnx.Rngs#

val_emb_dim: int#

class gensbi.models.TarFlow(params)[source]#

Bases: flax.nnx.Module

Transformer autoregressive normalizing flow density model.

Stacks MetaBlock bijections with alternating token permutations on top of a tokenizer and an isotropic Gaussian base distribution. Supports both vector and image data, with optional input standardization.

Parameters:: params (TarFlowParams) – Architecture and initialization parameters.

_base_log_prob(z)[source]#

Parameters:: z (jax.Array)
Return type:: jax.Array

_ensure_batched(x)[source]#

Parameters:: x (jax.Array)
Return type:: jax.Array

log_prob(x, cond=None)[source]#

Compute the log-probability of data under the model.

Applies standardization, tokenizes the input, then runs each MetaBlock’s inverse() transform (data→noise direction), accumulating the log-absolute-determinant terms, and finally evaluates the base Gaussian log-probability.

Parameters:

x (Array) – Data samples of shape (B, *example_shape) or a single unbatched sample that will be promoted to a batch of one.
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 (flattened internally by the conditioner). Pass None for an unconditional model.

Returns:

Log-probabilities of shape (B,).

Return type:

Array

sample(key, cond=None, nsamples=None)[source]#

Draw samples from the model.

Samples noise from N(0, I), then applies each MetaBlock’s forward() transform (noise→data direction) in reverse block order, detokenizes the result, and applies the inverse standardization.

Parameters:

key (jax.random.PRNGKey) – Random key for noise sampling.
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 by the conditioner). If provided, nsamples is inferred from cond.shape[0].
nsamples (int or None, optional) – Number of samples to draw. Required when cond is None.

Returns:

Samples of shape (B, *example_shape).

Return type:

Array

set_standardization(mean, std)[source]#

Set the mean and standard deviation for input standardization.

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

Parameters:

mean (Array) – Mean broadcastable to example_shape.
std (Array) – Standard deviation broadcastable to example_shape.

Return type:

None

Raises:

ValueError – If the model was built with standardize=False.

F = 1#

T#

_standardize#

blocks#

cond_dim#

dim#

example_shape#

mean#

std#

tokenizer#

class gensbi.models.TarFlowParams[source]#

Architecture parameters for TarFlow.

modeled selects the tokenizer ("vector" or "image"); cond selects the conditioner ("bias", "vector", or "image"). Head sizing follows the Flux1 convention: specify head_dim and num_heads; total width channels = head_dim * num_heads is derived in __post_init__.

Parameters:

rngs (nnx.Rngs) – Flax RNG container passed to all sub-modules during construction.
dim (int or None, optional) – Feature dimension of each input vector. Required when modeled="vector". Default is None.
cond_dim (int, optional) – Dimensionality of the conditioning vector. Set to 0 for an unconditional model. Default is 0.
modeled (str, optional) – Tokenizer type: "vector" (1-D data) or "image" (spatial data). Default is "vector".
img_size (int or None, optional) – Spatial size (height = width) of the modeled image. Required when modeled="image". Default is None.
patch_size (int or None, optional) – Patch size for the image tokenizer. Required when modeled="image". Default is None.
img_channels (int, optional) – Number of channels in the modeled image. Default is 1.
cond (str, optional) – Conditioning strategy: "bias" (per-token additive bias via AdditiveBiasConditioner), "vector" (one condition token per modeled coordinate via VectorConditioner), or "image" (prefix tokens from an image via ImageConditioner). Default is "bias".
cond_img_size (int or None, optional) – Spatial size of the conditioning image. Required when cond="image". Default is None.
cond_patch_size (int or None, optional) – Patch size for the image conditioning tokenizer. Required when cond="image". Default is None.
cond_channels (int, optional) – Number of channels in the conditioning image. Default is 1.
head_dim (int, optional) – Dimension per attention head. Default is 16.
num_heads (int, optional) – Number of attention heads per block. Default is 4.
num_blocks (int, optional) – Number of MetaBlock layers. Default is 8.
layers_per_block (int, optional) – Number of AttentionBlock layers inside each MetaBlock. Default is 2.
block_size (int, optional) – Token grouping factor for the vector tokenizer. Default is 1.
permutation (str, optional) – Token permutation strategy per block: "flip" (alternate forward/reverse order) or "random" (independently sampled per block). Default is "flip".
standardize (bool, optional) – If True (default), apply mean/std standardization to inputs and outputs. Enables TarFlow.set_standardization().
zero_init (bool, optional) – If True (default), initialize proj_out weights to zero so each MetaBlock starts as the identity map.
use_softplus (bool, optional) – If True (default), use softplus for the affine scale (numerically stable, bounded tail). If False, use exp (legacy behavior).
soft_clip (float, optional) – Soft-clip magnitude applied via tanh to raw network outputs before splitting into (a, b). Default is 4.0.
use_rope (bool, optional) – If True, replace the learned per-token pos_embed for the modeled image tokens with 2D rotary position embeddings (VisionRotaryEmbedding). Prefix (condition) tokens keep their learned embeddings and sit at the identity rotation (zero angles). Requires modeled="image" and head_dim divisible by 4. head_dim >= 32 is recommended for image data (more rotary frequencies per axis) but not enforced. Default is False.
rope_theta (int, optional) – Frequency base for the rotary embedding. Only used when use_rope=True. Default is 10000.
param_dtype (DTypeLike, optional) – Dtype for all stored (master) kernel/bias/embedding parameters across the tokenizer, conditioner, and transformer blocks. Default is float32.
dtype (DTypeLike, optional) – Compute dtype knob threaded through the conditioners and MetaBlock/ AttentionBlock layers. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast. Hard-fp32 regardless of this knob: the softplus/soft_clip affine-scale path in MetaBlock._affine, log-det accumulation, the mean/std standardization buffers, and the KV-cache buffers used during sampling.

__post_init__()[source]#

block_size: int = 1#

cond: str = 'bias'#

cond_channels: int = 1#

cond_dim: int = 0#

cond_img_size: int | None = None#

cond_patch_size: int | None = None#

dim: int | None = None#

dtype: jax.typing.DTypeLike#

head_dim: int = 16#

img_channels: int = 1#

img_size: int | None = None#

layers_per_block: int = 2#

modeled: str = 'vector'#

num_blocks: int = 8#

num_heads: int = 4#

param_dtype: jax.typing.DTypeLike#

patch_size: int | None = None#

permutation: str = 'flip'#

rngs: flax.nnx.Rngs#

rope_theta: int = 10000#

soft_clip: float = 4.0#

standardize: bool = True#

use_rope: bool = False#

use_softplus: bool = True#

vec_channels: int = 1#

zero_init: bool = True#

class gensbi.models.UnconditionalWrapper(model)[source]#

Bases: gensbi.utils.model_wrapping.ModelWrapper

Wrapper for unconditional models to handle input expansion and calling convention.

Parameters:: model (The unconditional model instance to wrap.)

__call__(t, obs, obs_ids, **kwargs)[source]#

Call the wrapped model with expanded inputs.

Parameters:

t (Array) – Time steps.
obs (Array) – Observations.
obs_ids (Array) – Observation identifiers.
**kwargs (Additional keyword arguments passed to the model.)

Returns:

Model output.

Return type:

Array