Flax (NNX) Cheat Sheet

1Install & a ModuleNNX

pip install flax # brings in jax; add optax★
Install JAX for your accelerator separately (jax[cuda12]). from flax import nnx.
class MLP(nnx.Module): def __init__(self, din, dh, dout, *, rngs: nnx.Rngs): self.l1 = nnx.Linear(din, dh, rngs=rngs) self.l2 = nnx.Linear(dh, dout, rngs=rngs)★
Create sublayers in __init__ and store them as attributes — just like PyTorch. Pass an nnx.Rngs for init randomness.
def __call__(self, x): return self.l2(nnx.relu(self.l1(x)))★
Define the forward in __call__. No @nn.compact, no separate apply — it's a normal method.
model = MLP(2, 16, 3, rngs=nnx.Rngs(0)) y = model(x) # eager call, weights already initialized★
Instantiation runs initialization (unlike Linen's lazy init). Then call the model directly.

2Built-in layersnnx.*

nnx.Linear(din, dout, rngs=rngs) · nnx.Conv(...) · nnx.Embed(...)★
Core layers all take rngs=. Weight init via kernel_init=, bias_init=.
nnx.BatchNorm(dim, rngs=rngs) · nnx.LayerNorm(dim) · nnx.Dropout(0.1, rngs=rngs)★
Stateful/stochastic layers — they hold running stats / need RNG, handled by the module system.
nnx.MultiHeadAttention(...) · nnx.relu / gelu / softmax (via jax.nn)
Attention block plus the usual activations (nnx re-exports jax.nn functions).
nnx.Sequential(nnx.Linear(...), nnx.relu, nnx.Linear(...))
Quick sequential container; or write a custom Module for anything non-linear.

3Params, state, train/evalwhat's inside

self.w = nnx.Param(jax.random.normal(key, (din, dout)))★
Learnable weights are nnx.Param. Non-trainable state uses nnx.Variable subclasses (e.g. nnx.BatchStat).
nnx.state(model) · nnx.state(model, nnx.Param)★
Extract the model's state (all variables, or filtered by type) as a pytree — for saving, counting, or transforms.
model.train() · model.eval()★
Toggle Dropout/BatchNorm behavior — same idiom as PyTorch. Sets the deterministic/use-running-stats flags.
nnx.display(model) · nnx.tabulate(model, x)
Inspect structure & parameter shapes/counts.

4grad & Optimizeroptax-backed

import optax optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)★
nnx.Optimizer wraps an Optax optimizer and holds its state and a reference to the model.
loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)★
NNX's grad transforms differentiate w.r.t. the model's Params directly — grads come back shaped like the model.
optimizer.update(model, grads)★
Applies the Optax update in place to the model's params (reference semantics) — no reassigning a new params dict.
nnx.grad(f) · nnx.value_and_grad(f, has_aux=True)
has_aux to return extra values (metrics) alongside the gradient.

5The training stepput it together

@nnx.jit def train_step(model, optimizer, x, y): def loss_fn(m): return ((m(x) - y)**2).mean() loss, grads = nnx.value_and_grad(loss_fn)(model) optimizer.update(model, grads) return loss★
The canonical step. model & optimizer are mutated in place; @nnx.jit handles their state under the hood.
for x, y in dataloader: loss = train_step(model, optimizer, x, y)★
A plain Python loop — because NNX is stateful, there's no carrying params/opt-state through the loop yourself.
nnx.MultiMetric(loss=nnx.metrics.Average(), acc=...)
Optional metric containers that accumulate across the epoch.

6nnx transforms (jit / vmap)JAX-aware

@nnx.jit · @nnx.vmap · @nnx.grad · @nnx.scan★
Use the nnx.* versions of JAX transforms on functions that take Modules — they understand NNX's reference state. Don't wrap a Module with bare jax.jit.
nnx.vmap(create_model, axis_size=8) # ensemble of 8 models
Vectorize over model instances (ensembles) or batch dims, state-aware.
nnx.jit(train_step, static_argnums=...)
Same knobs as jax.jit (static args, donation) — but state handling is automatic.

7split / merge (pure JAX)functional core

graphdef, state = nnx.split(model)★
Decompose a stateful Module into a static graphdef + a pytree state — both are pure JAX values you can pass through any jax.jit/grad.
model = nnx.merge(graphdef, state)★
Reconstruct the Module inside a functional context, run it, then split again. This is the escape hatch to raw JAX.
nnx.update(model, new_state) · nnx.state(model)
Push a modified state back onto a model, or read it out — for custom optimizers / manual manipulation.
graphdef, params, batch_stats = nnx.split(model, nnx.Param, nnx.BatchStat)
Split state into typed groups when you need to treat params and running stats differently.

8Checkpointingsave weights

import orbax.checkpoint as ocp _, state = nnx.split(model) ckpt = ocp.StandardCheckpointer(); ckpt.save(path, state)★
Save the state pytree with Orbax (the JAX-ecosystem standard). Params are just arrays.
restored = ckpt.restore(path, target=state) model = nnx.merge(graphdef, restored)★
Restore into an abstract state, then merge with a freshly-built model's graphdef.
nnx.state(model, ...) # count params: sum(x.size for x in jax.tree.leaves(...))
Because state is a pytree, all the JAX tree utilities apply.

9Linen (legacy) & migrationthe old API

import flax.linen as nn class MLP(nn.Module): @nn.compact def __call__(self, x): return nn.Dense(3)(x)
Linen is functional: modules are stateless templates, params live outside.
params = model.init(rng, x) y = model.apply(params, x)
The Linen loop: init to get params, apply to run — you thread params everywhere yourself.
# NNX ↔ Linen bridges exist for gradual migration
Linen isn't deprecated, but new code should prefer NNX. See the docs' evolution/migration guide.

!Common gotchasread before shipping

use nnx.jit, not jax.jit, on Modules
Bare JAX transforms don't understand NNX reference state. Use nnx.jit/grad/vmap, or split to pure state first.
pass rngs to stochastic layers
Dropout/init need an nnx.Rngs. Forgetting it (or reusing a key) breaks randomness — RNG is explicit, like all JAX.
train() vs eval()
Set the mode or Dropout stays on / BatchNorm uses batch stats at inference — a silent accuracy bug.
NNX ≠ Linen APIs
Don't mix tutorials: Linen uses @nn.compact+init/apply; NNX uses plain __init__/__call__ and in-place updates.

Flax cheat sheet