gensbi.models.tarflow.blocks

Transformer blocks for the transformer flow.

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

Attributes

INV_SOFTPLUS_1

Classes

AttentionBlock	Pre-norm residual block combining causal self-attention and an MLP.
MetaBlock	One exact autoregressive bijection over a token sequence.

Module Contents

class gensbi.models.tarflow.blocks.AttentionBlock(channels, num_heads, expansion, rngs, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

Pre-norm residual block combining causal self-attention and an MLP.

LayerNorm is applied over the channel axis only (not across tokens), so no future-token information leaks into earlier positions.

Parameters:

• channels (int) – Token embedding width. Must be divisible by num_heads.

• num_heads (int) – Number of attention heads.

• expansion (int) – MLP hidden-size multiplier; the MLP has channels * expansion neurons in its hidden layer.

• rngs (nnx.Rngs) – Flax RNG container used to initialize all sub-layers.

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

• dtype (DTypeLike, optional) – Compute dtype forwarded to each Linear. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast. norm1/norm2 are fp32 islands and always run at float32 regardless of this knob; their output self-heals when it feeds the following compute-dtype Linear (promote_dtype downcasts it there).

__call__(x, mask=None, freqs_cis=None, return_kv=False)

Apply pre-norm residual self-attention followed by a residual MLP.

Parameters:

• x (Array) – Token sequence of shape (B, T, C).

• mask (Array or None, optional) – Attention mask broadcastable to (1, 1, T, T); None means standard causal.

• freqs_cis (Array or None, optional) – Rotary angles (T, head_dim) applied to q/k before attention.

• return_kv (bool, optional) – If True, also return the unrotated (k, v) — the KV cache stores unrotated keys (the cached path re-rotates the full cache each step, as the reference does). Default False.

Returns:

out of shape (B, T, C), or (out, k, v) when return_kv=True.

Return type:

Array or tuple

_finish(x, attn)

Residual attention output + residual MLP (shared tail).

Parameters:

• x (jax.Array)

• attn (jax.Array)

Return type:

jax.Array

_qkv(x)

Project to unrotated (q, k, v), each of shape (B, S, nh, hd).

Parameters:

x (jax.Array)

decode(x_new, k_cache, v_cache, index, freqs_cis=None)

Single-token decode step against a preallocated KV cache.

The cache stores unrotated k (reference behavior: rope is applied after the cache read, re-rotating the whole prefix each step). Slots beyond index are zero-filled and masked out of the attention.

Parameters:

• x_new (Array) – New token, shape (B, 1, C).

• k_cache (Array) – Caches of shape (B, S, nh, hd) with S total slots.

• v_cache (Array) – Caches of shape (B, S, nh, hd) with S total slots.

• index (int or traced scalar) – Slot to write; attention sees slots <= index.

• freqs_cis (Array or None, optional) – Rotary angles (S, head_dim) for all slots; the new token’s q uses row index.

Returns:

(out, k_cache, v_cache) with out of shape (B, 1, C).

Return type:

tuple

head_dim

mlp_in

mlp_out

norm1

norm2

num_heads

proj

qkv

class gensbi.models.tarflow.blocks.MetaBlock(F, channels, T, perm, conditioner, num_layers, num_heads, expansion, rngs, zero_init=True, use_softplus=True, soft_clip=4.0, rope=None, grid=None, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

One exact autoregressive bijection over a token sequence.

Implements the Bijection direction contract: inverse() maps data to noise (density-evaluation direction) via a single parallel affine pass with a triangular Jacobian; forward() maps noise to data (sampling direction) via a sequential causal scan that re-runs the attention pass at each token position.

The affine scale is computed via softplus by default (bounded tail, identity at zero-init) or via exp when use_softplus=False (legacy).

Parameters:

• F (int) – Feature dimension per token (number of input channels per token).

• channels (int) – Internal embedding width for the transformer blocks.

• T (int) – Number of tokens in the sequence.

• perm (Array) – Token permutation applied before the affine transform. The inverse permutation is derived internally via argsort.

• conditioner (AdditiveBiasConditioner or VectorConditioner or ImageConditioner) – Module that provides (bias, prefix) conditioning signals via its embed method.

• num_layers (int) – Number of AttentionBlock layers stacked inside this block.

• num_heads (int) – Number of attention heads passed to each AttentionBlock.

• expansion (int) – MLP expansion factor passed to each AttentionBlock.

• rngs (nnx.Rngs) – Flax RNG container used to initialize all sub-layers.

• zero_init (bool, optional) – If True (default), initialize the output projection proj_out to zero so the block starts as the identity map.

• use_softplus (bool, optional) – If True (default), use softplus to compute 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. Set to 0 to disable clipping.

• rope (VisionRotaryEmbedding or None, optional) – If given, 2D rotary position embeddings replace the learned pos_embed (which is then set to None). Positions are laid out with the M prefix (conditioning) slots at the identity rotation (zero angles) followed by the T image slots on the normalized grid (raster order, not permuted by perm). Default is None (learned pos_embed, unchanged behavior).

• grid (tuple of int or None, optional) – (h, w) patch-grid shape used to build rope positions when rope is given; required in that case. Default is None.

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

• dtype (DTypeLike, optional) – Compute dtype forwarded to proj_in/proj_out/AttentionBlock layers. Default is float32, matching param_dtype, so with default arguments this is a bit-identical no-op cast. The softplus/soft_clip affine-scale computation in _affine() stays unconditionally fp32 regardless of this knob.

_affine(a)

Map raw log-scale a -> (scale, inv_scale, log_scale) in float32.

scale plays the role of exp(a) (“1/sigma”): inverse multiplies by inv_scale, forward multiplies by scale, logdet sums log_scale. softplus mode bounds the positive-scale tail and its gradient; the INV_SOFTPLUS_1 offset makes it the identity at a == 0.

Parameters:

a (jax.Array)

_embed_cond(cond)

Condition-only signals (bias, prefix, mask) for _params_core().

These depend on cond but not on the modeled tokens, so the sampling scan computes them once instead of re-running the (potentially expensive, e.g. image-patchify) conditioner and rebuilding the prefix mask at every token step.

Parameters:

cond (jax.Array | None)

_forward_reference(z, cond=None)

Map noise to data via full recompute (reference path for the KV cache).

forward is the production path (KV-cached); this method is its correctness oracle, retained for the equivalence test suite — do not delete it.

Sequentially scans over token positions via jax.lax.scan, re-running the causal attention pass at each step so that token i’s parameters are conditioned on already-generated tokens 0, …, i-1 (mirrors MaskedAutoregressive.forward).

Parameters:

• z (Array) – Noise-space token sequence of shape (B, T, F) or a flat array that will be reshaped to (B, T, F).

• cond (Array or None, optional) – Conditioning input, or None for an unconditional transform.

Returns:

• x (Array) – Data-space output of shape (B, T, F).

• logabsdet (Array) – Log absolute determinant of the Jacobian of the forward map, shape (B,). Equal to +Σ log_scale over token and feature dimensions.

_params(x_perm, cond)

(a, b) for the permuted tokens (single-shot; embeds the condition).

Parameters:

• x_perm (jax.Array)

• cond (jax.Array | None)

_params_core(x_perm, bias, prefix, mask)

(a, b) for the permuted tokens given precomputed conditioning.

SOS input-shift makes token i’s params depend only on tokens < i (and the condition). bias/prefix/mask come from _embed_cond().

Parameters:

x_perm (jax.Array)

_prefix_mask(M, T)

Prefix-LM mask over [prefix(M); modeled(T)]: modeled is causal and sees all prefix; prefix is bidirectional among itself and never sees modeled (cond→x blocked).

Parameters:

• M (int)

• T (int)

Return type:

jax.Array

forward(z, cond=None)

Map noise to data (the sampling direction), KV-cached.

Prefills the per-layer caches with the condition prefix (one parallel pass under the bidirectional prefix mask, matching the training-path mask rows), then scans over token positions decoding a single token per step against the caches. Verified equivalent to _forward_reference() (full recompute) by the test suite.

Parameters:

• z (Array) – Noise-space token sequence of shape (B, T, F) or a flat array that will be reshaped to (B, T, F).

• cond (Array or None, optional) – Conditioning input, or None for an unconditional transform.

Returns:

• x (Array) – Data-space output of shape (B, T, F).

• logabsdet (Array) – Log absolute determinant of the Jacobian of the forward map, shape (B,). Equal to +Σ log_scale over token and feature dimensions.

inverse(x, cond=None)

Map data to noise (the density-evaluation direction).

Applies a per-token parallel affine transform z = (x − b) · inv_scale after permuting tokens. Token i’s parameters (a, b) depend only on tokens at positions < i (causal attention on a shift-by-one input), giving a triangular Jacobian computable in a single forward pass.

Parameters:

• x (Array) – Data-space token sequence of shape (B, T, F) or a flat array that will be reshaped to (B, T, F).

• cond (Array or None, optional) – Conditioning input, or None for an unconditional transform.

Returns:

• z (Array) – Noise-space output of shape (B, T, F).

• logabsdet (Array) – Log absolute determinant of the Jacobian of the inverse map, shape (B,). Equal to -Σ log_scale over token and feature dimensions.

F

T

attn_blocks

conditioner

dtype

inv_perm

perm

proj_in

proj_out

soft_clip = 4.0

sos_embed

use_softplus = True

gensbi.models.tarflow.blocks.INV_SOFTPLUS_1 = 0.541324854612918