gensbi.models.tarflow#

Transformer autoregressive normalizing-flow density model.

Exports TarFlow and its configuration dataclass TarFlowParams for building autoregressive normalizing flows based on a stack of causal transformer blocks (MetaBlocks).

Submodules#

Classes#

`TarFlow`	Transformer autoregressive normalizing flow density model.
`TarFlowParams`	Architecture parameters for `TarFlow`.

Package Contents#

class gensbi.models.tarflow.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.tarflow.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#