Pytrees

Take any nested structure you'd build to hold a model's weights: a dict of dicts, a list of arrays, maybe a tuple mixed in. JAX has a name for exactly that shape, a pytree, any nesting of dicts, lists, tuples, and registered node types, bottoming out in leaves like arrays and scalars. jax.tree.flatten takes one of these and splits it into two things: a flat list of leaves, and a treedef that remembers how to put them back.

This is not a side detail. It is the calling convention every transformation in JAX shares. Call grad on a loss that takes a dict of arrays, and the gradient that comes back is a dict with the exact same keys and shapes, same treedef as the params argument, by construction. jit and vmap do the same thing on the way in, flattening whatever you pass to leaves before tracing touches it. That is why reaching for params = {"dense": {...}, "output": {...}} instead of some bespoke class is idiomatic JAX, not a beginner's habit you'll grow out of.

Every transformation sees your model the same way: leaves on a tree.

One habit from Python trips people here. None looks like it should be a leaf, it's a value, after all, but JAX treats it as an empty subtree: it vanishes from the leaves list entirely. And a list is not a tuple, structurally, even when they hold the same values in the same order. Zip a dict of lists against a dict of tuples with tree.map and you don't get silent coercion. You get a structure mismatch error, because those are two different treedefs pretending to look alike.

Once every model is leaves and a treedef, updating it stops needing an architecture-specific loop. jax.tree.map(f, tree, *rest) walks every tree in lockstep, structure required to match, and applies f leafwise. Wire the SGD update as params - 0.1 * grads and it is one line, unchanged whether the model has three parameters or three hundred.

The rest of the kit rounds out the same idea. jax.tree.leaves gets you the flat list without the treedef, jax.tree.structure gets you the treedef without the leaves, and jax.tree.unflatten puts a new set of leaves back into an old structure. Together with flatten and map, that is the whole surface most training code needs.

import jax
import jax.numpy as jnp

params = {"dense": {"w": jnp.ones((2, 2)), "b": jnp.zeros(2)}, "scale": jnp.ones(())}
grads = jax.tree.map(jnp.ones_like, params)
new_params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
print(jax.tree.structure(new_params))   # same treedef as params, by construction

A plain dataclass is not automatically a pytree. JAX doesn't know which of its fields are arrays that should flow through transformations and which are configuration that should stay fixed, so left alone, jax.grad or jax.jit would just refuse it, or worse, treat the whole object as one opaque leaf. @jax.tree_util.register_dataclass tells JAX how to read it: which fields are data, which are static, declared right on the field with dataclasses.field(metadata=dict(static=True)).

That decorator covers the common case, a dataclass that's mostly arrays with a few static settings. For a container you want full control over, custom flatten and unflatten logic, not just field-by-field mapping, register_pytree_node is the lower-level tool underneath it. Either way, the lesson is the same: a custom class doesn't opt out of the pytree system by existing. It opts in the moment you register it, and from then on grad, jit, and vmap walk through it exactly as they would a dict.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@jax.tree_util.register_dataclass
@dataclass
class TrainState:
    step: jax.Array
    params: dict

state = TrainState(step=jnp.zeros((), jnp.int32), params={"w": jnp.ones(3)})
bumped = jax.tree.map(lambda x: x + 1, state)
print(bumped.step)   # 1: the dataclass is transparent to every transform