gensbi.models.tarflow.conditioners

Conditioning seams for the transformer flow.

Adapted from apple/ml-tarflow (TarFlow); see models/tarflow/LICENSE.apple. Prefix-concatenation conditioning adapted from apple/ml-starflow (STARFlow); see models/tarflow/LICENSE.starflow.

AdditiveBiasConditioner is the continuous analog of TarFlow’s class_embed: an MLP embeds the condition to a channels-vector that is broadcast-added to every token. The signal depends only on the condition (constant w.r.t. the modeled variable), so it shifts the affine params without breaking the triangular Jacobian. A plain 2-layer MLP is used (not MLPEmbedder, whose hidden_dim % in_dim == 0 constraint does not fit arbitrary cond_dim).

Classes

AdditiveBiasConditioner	Embed a vector condition as a per-token additive bias.
ImageConditioner	Embed an image condition as prefix tokens prepended to the sequence.
VectorConditioner	Embed a vector condition as one prefix token per coordinate.

Module Contents

class gensbi.models.tarflow.conditioners.AdditiveBiasConditioner(cond_dim, channels, rngs, cond_channels=1, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

Embed a vector condition as a per-token additive bias.

A two-layer MLP maps the condition to a channels-dimensional vector that is broadcast-added to every token in the sequence. When cond_dim == 0 the conditioner is unconditional and embed() returns (None, None).

Parameters:

• cond_dim (int) – Condition dimensionality. Set to 0 for an unconditional model.

• channels (int) – Output channel width matching the transformer embedding dimension.

• rngs (nnx.Rngs) – Flax RNG container for linear layer initialization.

• cond_channels (int, optional) – Number of channels in the conditioning input (B, cond_dim, C_cond). Default is 1. The input linear layer is widened to accept cond_dim * cond_channels features, folding the channel axis before the MLP (same flattening performed in embed()).

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

• dtype (DTypeLike, optional) – Compute dtype forwarded to the MLP layers. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast.

embed(cond)

Embed the condition into a per-token additive bias.

Parameters:

cond (Array or None) – Condition vector of shape (B, cond_dim) or (B, cond_dim, C_cond), or None when the model is unconditional (cond_dim == 0).

Returns:

• bias (Array or None) – Per-token additive bias of shape (B, channels), or None when cond_dim == 0.

• prefix (None) – This conditioner does not produce prefix tokens; always None.

Raises:

ValueError – If cond is None when cond_dim > 0.

cond_channels = 1

cond_dim

class gensbi.models.tarflow.conditioners.ImageConditioner(cond_channels, patch_size, channels, num_tokens, rngs, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

Embed an image condition as prefix tokens prepended to the sequence.

Patchifies a spatial image (B, H, W, C) into M = (H / patch_size) * (W / patch_size) flat patch vectors, projects each patch to channels dimensions, and adds learned positional embeddings.

Parameters:

• cond_channels (int) – Number of channels in the conditioning image.

• patch_size (int) – Spatial size of each square patch (height and width in pixels).

• channels (int) – Output channel width matching the transformer embedding dimension.

• num_tokens (int) – Number of prefix tokens; must equal (H / patch_size) * (W / patch_size).

• rngs (nnx.Rngs) – Flax RNG container for projection layer and positional embedding initialization.

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

• dtype (DTypeLike, optional) – Compute dtype forwarded to proj. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast.

embed(cond)

Patchify an image condition and embed it as prefix tokens.

Parameters:

cond (Array or None) – Image condition of shape (B, H, W, C).

Returns:

• bias (None) – This conditioner does not produce a per-token additive bias; always None.

• prefix (Array) – Prefix token sequence of shape (B, num_tokens, channels) with learned positional embeddings added.

Raises:

ValueError – If cond is None.

M

channels

patch_size

pos

proj

class gensbi.models.tarflow.conditioners.VectorConditioner(cond_dim, cond_channels, channels, rngs, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

Embed a vector condition as one prefix token per coordinate.

Each of the cond_dim coordinates is a token of C_cond channels; a shared Linear(cond_channels, channels) projects each to the transformer width, plus per-coordinate positional embeddings. Produces M = cond_dim prefix tokens (no flatten).

Parameters:

• cond_dim (int) – Condition dimensionality (number of coordinates / prefix tokens).

• cond_channels (int) – Number of channels per coordinate in the input condition (B, cond_dim, cond_channels).

• channels (int) – Output channel width matching the transformer embedding dimension.

• rngs (nnx.Rngs) – Flax RNG container for linear layer and positional embedding initialization.

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

• dtype (DTypeLike, optional) – Compute dtype forwarded to proj. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast.

embed(cond)

Embed the condition into per-coordinate prefix tokens.

Parameters:

cond (Array or None) – Condition array of shape (B, cond_dim, C_cond).

Returns:

• bias (None) – This conditioner does not produce a per-token additive bias; always None.

• prefix (Array) – Prefix token sequence of shape (B, cond_dim, channels) with learned positional embeddings added.

Raises:

ValueError – If cond is None.

M

channels

cond_channels

cond_dim

pos

proj