jax-tarflow

name	jax-tarflow
description	Use this skill when constructing, sampling, training, saving, loading, or documenting the JAX plus Equinox port of TarFlow, the autoregressive transformer normalizing flow from Apple's ml-tarflow paper.

JAX TarFlow

This skill documents a self-contained JAX/Equinox port of the TarFlow model from the paper "Normalizing Flows are Capable Generative Models" (Zhai et al., 2024, arXiv:2412.06329). The PyTorch reference lives in $ML_TARFLOW/transformer_flow.py, where ML_TARFLOW points at a local clone of Apple's ml-tarflow. The skill vendors a fresh port under jax_tarflow/ and ships four runnable examples plus a five-test suite that includes a PyTorch parity check.

The port covers the paper-default feature set: NVP affine coupling layers, class-conditional generation with a -1 sentinel for the unconditional path, and KV-cached autoregressive sampling with classifier-free guidance. It is fp32 only; mixed precision, the VP-mode learned prior variance, and multi-device sharding are explicitly out of scope.

• construct.py builds a tiny model and verifies forward + sample shapes.
• generate.py runs unconditional, class-conditional, and CFG sampling.
• train.py trains on MNIST for a configurable number of steps.
• save_load.py saves, lists, and reloads a model with a parity assertion.
• tests/ contains invertibility, log-determinant, causal-mask, KV-cache parity, and PyTorch parity tests.

Clone Apple's ml-tarflow repo somewhere local, then set ML_TARFLOW to point at that clone. Install the JAX dependencies into the ml-tarflow uv environment via the jax extra exposed in pyproject.toml, then run the examples or tests.

git clone https://github.com/apple/ml-tarflow ~/ml-tarflow
export ML_TARFLOW=~/ml-tarflow

uv sync --extra jax
uv run python ~/.claude/skills/jax-tarflow/construct.py
uv run python ~/.claude/skills/jax-tarflow/generate.py
uv run python ~/.claude/skills/jax-tarflow/train.py --steps 100
uv run python ~/.claude/skills/jax-tarflow/save_load.py
uv run pytest ~/.claude/skills/jax-tarflow/tests

Module Map

The top-level config and model live in jax_tarflow.config and jax_tarflow.model.

• jax_tarflow.config.TarflowConfig is the frozen dataclass that records every architectural hyperparameter.
• jax_tarflow.attention holds MLP, Attention, AttentionBlock, and the KVCache dataclass used by autoregressive sampling.
• jax_tarflow.meta_block holds PermutationIdentity, PermutationFlip, and the MetaBlock that pairs a permutation with a stack of AttentionBlocks.
• jax_tarflow.model holds patchify, unpatchify, and the top-level Model with __call__ for training, get_loss, and reverse for sampling.
• jax_tarflow.weight_port.port_weights_torch_to_jax copies parameters from a transformer_flow.Model into a JAX skeleton. It is used by the PyTorch parity test and is the canonical reference for tensor name mappings.

Construct

Model.init(cfg, key) builds the full model. Equinox stores parameters as a pytree of arrays; static fields like cfg, head_dim, and permutation type are marked eqx.field(static=True). The default head_dim is 64 to match the PyTorch reference, so channels must be divisible by 64. For tiny smoke tests use head_dim=16 with channels=32.

cfg = TarflowConfig(
    in_channels=1, img_size=8, patch_size=2,
    channels=32, num_blocks=2, layers_per_block=2,
    num_classes=4, head_dim=16,
)
model = Model.init(cfg, jax.random.PRNGKey(0))

See construct.py for a runnable version with shape assertions on both the forward pass and a sampling call.

Generate

generate.py shows the three sampling call signatures. The latents noise are drawn outside the model, since Model.reverse is deterministic given its input noise.

noise = jax.random.normal(key, (batch, cfg.num_patches, cfg.patch_channels))

unconditional = model.reverse(noise, y=None)
conditional = model.reverse(noise, y=jnp.array([0, 1, 2, 3]))
cfg_sample = model.reverse(
    noise, y=jnp.array([0, 1, 2, 3]),
    guidance=2.3, attn_temp=1.0, annealed_guidance=False,
)

Class conditioning uses a -1 sentinel: any batch entry where y == -1 is routed through the mean class embedding, which is also the unconditional path.

Train

The training step is a plain eqx.filter_value_and_grad over Model.get_loss on the forward output, with an optax.adamw optimizer and eqx.apply_updates. JIT lives on train_step, the outermost call that covers the full computation, per the project's JIT discipline.

@eqx.filter_jit
def train_step(model, opt_state, x, y, noise_key, drop_key, optim, noise_std, drop_label):
    x = x + noise_std * jax.random.normal(noise_key, x.shape)
    drop_mask = jax.random.bernoulli(drop_key, p=drop_label, shape=y.shape)
    y_dropped = jnp.where(drop_mask, -1, y)
    loss, grads = eqx.filter_value_and_grad(loss_fn)(model, x, y_dropped)
    updates, opt_state = optim.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss

CFG-dropout training replaces each row's label with -1 at a drop_label rate, which is the same vectorized trick PyTorch tarflow uses (mask = rand < p; y = (1 - mask) * y - mask).

The MNIST recipe matches $ML_TARFLOW/train_local.ipynb: patch_size=4, channels=128, num_blocks=4, layers_per_block=4, AdamW betas=(0.9, 0.95), lr=2e-3, weight_decay=1e-4, noise_std=0.1 gaussian dequantization. The script ships without a learning-rate schedule; for longer runs plug in optax.warmup_cosine_decay_schedule.

Sample

The autoregressive reverse walks one patch at a time, predicting the affine update for patch i + 1 from patch i. To keep this lax.scan-friendly under JIT, the attention layers maintain a pre-allocated KV cache.

The cache is a plain pytree (KVCache with k, v, and write_idx fields) threaded through the scan carry, written one slot at a time with jax.lax.dynamic_update_slice_in_dim. This is functionally identical to the eqx.nn.StateIndex pattern documented for stateful Equinox modules, but it avoids the eqx.nn.make_with_state ceremony since the cache is local to a single reverse call and is discarded on return.

# inside Attention.step
new_k = jax.lax.dynamic_update_slice_in_dim(cache.k, k_new, cache.write_idx, axis=1)
new_v = jax.lax.dynamic_update_slice_in_dim(cache.v, v_new, cache.write_idx, axis=1)
new_write_idx = cache.write_idx + 1
valid = jnp.arange(T) < new_write_idx
attn = jnp.einsum("bmhd,bnhd->bhmn", q * sqrt_scale, new_k * sqrt_scale) / temp
attn = jnp.where(valid[None, None, None, :], attn, -jnp.inf)
attn = jax.nn.softmax(attn, axis=-1)

Classifier-free guidance is implemented by doubling the batch dimension: the conditional batch and the unconditional batch run through one scan together, share the same noise trajectory, and combine their affine outputs per guide_what (subset of 'a' for scale, 'b' for shift) at every step. annealed_guidance=True linearly ramps the guidance scale from 0 to guidance over the patch sequence. attn_temp divides the attention logits during sampling only; the training path is unaffected.

When guidance == 0 the batch is not doubled, so unconditional and plain-conditional sampling have no CFG overhead.

Save And Load

save_load.py mirrors the jax-stylegan3 convention: each saved model is one directory under MODEL_ROOT = Path("pretrained_models") with weights.eqx and metadata.json. The metadata serializes the TarflowConfig so the skeleton can be rebuilt before deserializing leaves.

eqx.tree_serialise_leaves(model_dir / "weights.eqx", model)
metadata = {
    "name": name,
    "description": description,
    "saved_at": datetime.now(timezone.utc).isoformat(),
    "step": int(step),
    "loss": float(loss),
    "num_parameters": count_array_parameters(model),
    "config": asdict(cfg),
}

To load, reconstruct the skeleton from the saved config and then deserialize.

metadata = read_metadata(model_dir)
cfg = TarflowConfig(**metadata["config"])
skeleton = Model.init(cfg, key)
model = eqx.tree_deserialise_leaves(model_dir / "weights.eqx", skeleton)

No pretrained JAX checkpoint ships with this skill. The example saves a freshly initialized tiny model and asserts an exact output roundtrip after reload.

Tests

Five tests under tests/ exercise the port from two complementary angles.

• test_logdet.py checks that the per-block log-determinant equals slogdet(jax.jacobian(forward)) on a tiny single-block model, accounting for the per-dimension normalization in the loss convention.
• test_causal_mask.py perturbs the last patch and asserts earlier latents are unchanged.
• test_invertibility.py runs forward then reverse and asserts the input is recovered to fp32 precision (NVP is exactly invertible).
• test_kv_cache_parity.py compares the cached lax.scan reverse against a naive full-attention reverse at every step.
• test_torch_parity.py constructs a tiny PyTorch tarflow, copies its weights into a JAX skeleton via port_weights_torch_to_jax, and asserts forward parity within fp32 drift on both conditional and unconditional inputs.

Run the suite with uv run pytest tests/.

Notes For Agents

• Keep JAX code in this style: equinox modules with parameters as pytree leaves, no eqx.filter_jit on inner methods. Only train_step and example-level generate_one should be JIT'd, so XLA compiles the whole graph at once.
• eqx.filter_vmap is preferred over jax.vmap only when filtering parameters out of the vmap is required; the Attention here uses plain jax.vmap on already-pure functions.
• The -1 class sentinel is load-bearing in two places: the CFG-dropout training step in train.py and the CFG sampling path in meta_block.reverse. Do not introduce Optional[Array] for y in either hot loop, because vectorized dropout depends on the sentinel.
• The PyTorch reference uses head_dim=64 baked into MetaBlock.__init__. For parity tests against the upstream Model class, use channels=64.
• Sampling is deterministic given the noise input. There is no PRNG key argument to Model.reverse; randomness lives entirely in the caller.
• The shipped port covers NVP only. Adding VP mode requires reintroducing the per-block var buffer with the update_prior exponential moving average from transformer_flow.Model.