State

Every stateful thing in a JAX program, an optimizer's momentum, an RNG stream, a batch-norm running mean, a step counter, rides the same pattern: step(state, x) -> (new_state, y). State goes in as an argument. A new state comes out as a return value. Nothing gets mutated in place, and nothing lives outside the function's inputs and outputs.

State is an argument and a return value. Nothing else.

This isn't a style preference; it's forced by how tracing works. A traced function that mutated some global object would bake in whatever value that global held at trace time, and every later call would silently drop the mutation, because the recording never captured it happening. Threading state explicitly is the only way a function's behavior stays honest under jit and grad.

optax.adam(1e-3) isn't an object that quietly tracks its own momentum the way an optimizer class in a mutation-based framework would. It's a GradientTransformation: a pair of functions. init(params) builds the optimizer's state, itself a pytree of moment estimates the same shape as your params. update(grads, state, params) takes gradients and the current state and returns (updates, new_state). apply_updates(params, updates) folds those updates into your params. Every piece of state Adam needs is a value you're handed back, not something hidden inside an object.

Composing optimizers works because every piece shares that same shape. optax.chain(clip_by_global_norm, adam, ...) links transformations end to end, and each link, whatever it does, exposes the same init and update pair. Add gradient clipping or weight decay to a training loop and nothing about how you call the optimizer changes; only the chain grows by one link.

import jax.numpy as jnp
import optax

opt = optax.adam(1e-3)
params = {"w": jnp.ones(3)}
opt_state = opt.init(params)

grads = {"w": jnp.full(3, 0.5)}
updates, opt_state = opt.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
print(params["w"])   # moved against the gradient, Adam-scaled

NNX modules look, at first glance, like ordinary Python classes: you build a Linear layer, assign it to self.up, and mutate attributes in __init__ the way you would anywhere else. That's deliberate. NNX wants module code to read like normal Python, mutation included, while everything underneath still obeys the functional pattern from section one.

The bridge is nnx.split(model), which pulls a module apart into a graphdef, the static structure, what layers exist, how they connect, and a state, a pytree of the actual arrays, the part that crosses jit and grad boundaries cleanly. nnx.merge puts the two back together into a live module. Flax's older API, Linen, skips this split step and has you thread params explicitly through every call instead; NNX is the same underlying model with the bookkeeping automated.

None of this is an exemption. NNX's whole value is ergonomics laid on top of the same functional core chapter one and this chapter both describe, not an escape hatch from it. Every array NNX manages still moves through jit, grad, and vmap as a pytree leaf; the module syntax just spares you from writing the flatten and unflatten by hand.

import jax
from flax import nnx

class Mlp(nnx.Module):
    def __init__(self, din, dhidden, rngs: nnx.Rngs):
        self.up = nnx.Linear(din, dhidden, rngs=rngs)
        self.down = nnx.Linear(dhidden, din, rngs=rngs)

    def __call__(self, x):
        return self.down(jax.nn.relu(self.up(x)))

model = Mlp(4, 16, nnx.Rngs(0))
graphdef, state = nnx.split(model)   # static structure, pytree of arrays

Saving a model means saving a pytree, and Orbax is the library built for exactly that: it saves and restores arbitrary pytrees, asynchronously so training doesn't stall on disk I/O, and with sharding awareness baked in, so a checkpoint written from a distributed array, chapter ten's world, restores correctly across whatever device layout reads it back.

The other place host and device meet is callbacks. jax.experimental.io_callback and jax.pure_callback let a traced function reach out and run ordinary Python on the host: writing a metric to a log, reading from a file, anything the compiled program itself can't do. They come with real costs, ordering guarantees to reason about, a performance hit every time the device has to pause and wait on the host, so the honest use case is something occasional, like logging metrics, not something on the hot path of every step.