jax-patterns

name	jax-patterns
description	JAX, Flax, Optax, and Equinox patterns for ML training. Covers JIT, vmap, pmap, TPU usage, and functional model design. Use when building or debugging JAX-based training pipelines.

JAX Patterns

Framework Selection

Framework	Style	Strengths	Weaknesses	Use When
Flax (NNX)	Pythonic, mutable	Google-backed, large ecosystem, NNX is latest API	Linen API is legacy; NNX still maturing	Default for new projects, TPU training
Equinox	PyTorch-like, pytree-native	Clean API, composable, feels natural	Smaller ecosystem, fewer examples	Prefer PyTorch style, research code
Haiku	Functional transforms	DeepMind ecosystem, simple	Maintenance mode (use Flax NNX instead)	Legacy DeepMind codebases only

Default recommendation: Flax NNX for production/TPU work. Equinox for research where API ergonomics matter.

Core JAX Concepts

JIT Compilation

import jax
import jax.numpy as jnp
from functools import partial

# Basic JIT -- function must be pure (no side effects)
@jax.jit
def compute(x, y):
    return jnp.dot(x, y) + jnp.sin(x)

# Static arguments -- use for values that change compilation (shapes, dtypes, booleans)
@partial(jax.jit, static_argnums=(2,))
def forward(params, x, use_dropout: bool):
    ...

vmap -- Automatic Batching

def single_loss(params, x, y):
    pred = model.apply(params, x)
    return jnp.mean((pred - y) ** 2)

batched_loss = jax.vmap(single_loss, in_axes=(None, 0, 0))  # batch over x, y; not params
per_sample_grads = jax.vmap(jax.grad(single_loss), in_axes=(None, 0, 0))  # for DP

pmap -- Data Parallelism

@jax.pmap
def train_step(state, batch):
    def loss_fn(params):
        logits = state.apply_fn(params, batch["input"])
        return optax.softmax_cross_entropy_with_integer_labels(
            logits, batch["label"]
        ).mean()

    grads = jax.grad(loss_fn)(state.params)
    # All-reduce gradients across devices
    grads = jax.lax.pmean(grads, axis_name="devices")
    state = state.apply_gradients(grads=grads)
    return state

# Data must be sharded: (num_devices, per_device_batch, ...)
sharded_batch = jax.tree.map(
    lambda x: x.reshape(jax.local_device_count(), -1, *x.shape[1:]),
    batch,
)
state = train_step(state, sharded_batch)

Flax NNX Model Definition

from flax import nnx
import optax

class MLP(nnx.Module):
    def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
        self.linear1 = nnx.Linear(din, dhidden, rngs=rngs)
        self.bn = nnx.BatchNorm(dhidden, rngs=rngs)
        self.linear2 = nnx.Linear(dhidden, dout, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)

    def __call__(self, x):
        x = nnx.relu(self.bn(self.linear1(x)))
        x = self.dropout(x)
        return self.linear2(x)

model = MLP(784, 256, 10, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adamw(1e-3))

JIT-Compiled Training Step (Flax NNX)

@nnx.jit
def train_step(model, optimizer, batch):
    def loss_fn(model):
        logits = model(batch["image"])
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits, batch["label"]
        ).mean()
        return loss, {"logits": logits}

    grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)
    (loss, aux), grads = grad_fn(model)
    optimizer.update(grads)
    return loss, aux

Equinox Alternative

import equinox as eqx
import optax

class MLP(eqx.Module):
    layers: list

    def __init__(self, key):
        k1, k2 = jax.random.split(key)
        self.layers = [
            eqx.nn.Linear(784, 256, key=k1),
            eqx.nn.Linear(256, 10, key=k2),
        ]

    def __call__(self, x):
        x = jax.nn.relu(self.layers[0](x))
        return self.layers[1](x)

model = MLP(jax.random.PRNGKey(0))
optimizer = optax.adamw(1e-3)
opt_state = optimizer.init(eqx.filter(model, eqx.is_array))

@eqx.filter_jit
def train_step(model, opt_state, batch):
    @eqx.filter_value_and_grad
    def loss_fn(model):
        logits = jax.vmap(model)(batch["image"])
        return optax.softmax_cross_entropy_with_integer_labels(
            logits, batch["label"]
        ).mean()

    loss, grads = loss_fn(model)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss

Optax Optimizer Chains

import optax

# Composed optimizer: gradient clipping + AdamW + weight decay schedule
schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0, peak_value=1e-3,
    warmup_steps=1000, decay_steps=50000, end_value=1e-5,
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),       # gradient clipping
    optax.adamw(learning_rate=schedule, weight_decay=0.01),
)

# Separate treatment for different param groups
def partition_fn(params):
    """Label params for different optimizer treatment."""
    flat = flax.traverse_util.flatten_dict(params)
    labels = {k: "no_decay" if "bias" in k[-1] or "norm" in k[-1] else "decay"
              for k in flat}
    return flax.traverse_util.unflatten_dict(labels)

optimizer = optax.multi_transform({
    "decay": optax.adamw(1e-3, weight_decay=0.01),
    "no_decay": optax.adam(1e-3),
}, partition_fn)

Custom Loss with Auxiliary Data

def loss_fn(params, batch, rngs):
    logits = model.apply(params, batch["input"], rngs=rngs)
    loss = optax.softmax_cross_entropy_with_integer_labels(
        logits, batch["label"]
    ).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == batch["label"])
    # Return loss as primary, metrics as aux
    return loss, {"accuracy": accuracy, "logits": logits}

grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
(loss, metrics), grads = grad_fn(state.params, batch, {"dropout": rng})

TPU-Specific Patterns

print(jax.devices())  # [TpuDevice(id=0, ...), ...]

# TPU uses bfloat16 natively for matmuls. Store params in float32 for stability.
# Multi-host pmap (TPU pods) -- pmean handles all-reduce across hosts
@jax.pmap(axis_name="batch")
def train_step(state, batch):
    ...
    grads = jax.lax.pmean(grads, axis_name="batch")
    ...

Gotchas and Anti-Patterns

Functional Constraints

• No in-place mutation: x[0] = 1 fails. Use x = x.at[0].set(1)
• No side effects in JIT: print() only runs at trace time. Use jax.debug.print() instead
• Random numbers require explicit key management: key, subkey = jax.random.split(key)

Pytree Handling

• All JAX transforms operate on pytrees. Custom classes need pytree registration or a framework
• None values in pytrees are silently dropped -- can cause shape mismatches
• Use jax.tree.map (not manual loops) for transformations across nested structures

XLA Compilation Caching

• First JIT call is slow (XLA compilation). Same-shape subsequent calls are fast
• Shape changes trigger recompilation -- pad to fixed shapes when possible
• Set JAX_COMPILATION_CACHE_DIR to persist compiled programs across runs

TPU Topology

• pmap handles 1D parallelism. For 2D+ mesh, use jax.sharding + mesh_utils
• Data loading bottleneck: TPU VMs have limited CPU/RAM. Use tf.data or Grain
• TPU bfloat16 is default for matmuls -- don't force float16, it's slower on TPU

Common Mistakes

• Using np instead of jnp inside JIT -- silent fallback to CPU
• Forgetting axis_name in pmap when using pmean/psum
• Creating arrays inside JIT -- they become constants baked into the compiled graph
• Reusing PRNG keys without splitting -- produces correlated "random" numbers
• Python for loops in JIT unroll fully -- use jax.lax.scan or fori_loop instead