Pytrees

Setup#

pip install jax

Outputs verified with jax 0.6.2, CPU. (jax.tree.map is the current spelling; older code says jax.tree_util.tree_map — same function.)

Parameters as a tree, updated in one line#

import jax

params = {
    "layer1": {"w": jax.numpy.ones((2, 3)), "b": jax.numpy.zeros(3)},
    "layer2": {"w": jax.numpy.ones((3, 1)), "b": jax.numpy.zeros(1)},
}

# One call reaches every array in the nest, however deep
shapes = jax.tree.map(lambda leaf: leaf.shape, params)
print(shapes)

leaves = jax.tree.leaves(params)
print("number of leaves:", len(leaves))
print("total parameters:", sum(leaf.size for leaf in leaves))

# A gradient update touches every leaf without naming any of them
grads = jax.tree.map(lambda leaf: leaf * 0 + 0.5, params)   # pretend gradients
updated = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
print("layer1 w after update:", updated["layer1"]["w"][0])

{'layer1': {'b': (3,), 'w': (2, 3)}, 'layer2': {'b': (1,), 'w': (3, 1)}}
number of leaves: 4
total parameters: 13
layer1 w after update: [0.95 0.95 0.95]

The walkthrough#

jax.tree.map(f, tree) applies f to each leaf and rebuilds the same structure around the results. The shapes printout keeps the dict nesting — instantly readable as a model summary. This shape-audit one-liner is worth memorising; it is the first thing to print when any JAX code misbehaves.

jax.tree.leaves flattens to a plain list, structure discarded — right for reductions like counting, wrong for anything that must go back into the model. Its inverse pair is jax.tree.flatten / jax.tree.unflatten, which keep the structure ticket so you can rebuild.

The two-tree form is the training step. jax.tree.map(lambda p, g: p - 0.1 * g, params, grads) walks both trees in lockstep — leaf pairs meet, combine, and a new tree returns. Every JAX training loop has this line, or hides it inside Optax. Check the arithmetic in the output: 1.0 − 0.1 × 0.5 = 0.95.

Everything accepts pytrees, not only tree.map. grad differentiates functions of pytrees, jit compiles them, vmap batches them — arguments and returns can be arbitrary nests everywhere in JAX. This is why the library needs no Parameter class and no Module base: plain containers are the parameter format. Structures you register yourself (dataclasses) can join too.

Common mistakes#

Mismatched structures in multi-tree map. If grads has a leaf params lacks — or a list where a tuple should be — you get ValueError: ... tree structure mismatch (wording varies). Diagnose with jax.tree.structure(a) == jax.tree.structure(b). The subtle cousin: None values, which are treated as empty structure, not a leaf, and silently vanish from maps.

Forgetting a tuple is a tree too. Return (loss, params) from a function and tree operations will happily descend into it. If you meant "treat this whole thing as one value", the is_leaf= argument draws the boundary.

Loops over .items() doing what map does. Hand-walking nested dicts works until the nesting changes depth, then breaks. tree.map is depth-proof — code written for layer1/layer2 survives becoming block/layer/sublayer untouched.

Assuming an ordering. Dict leaves flatten in sorted-key order, not insertion order. Code that pairs tree.leaves(params) with tree.leaves(grads) by position works only because both sort the same way — rely on matched structure, not on remembered order.

Try it yourself#

Compute the largest absolute value across all parameters in updated — flatten with jax.tree.leaves, then one max. Then write clip = lambda t: jax.tree.map(lambda x: jax.numpy.clip(x, -0.9, 0.9), t) and confirm shapes survive: jax.tree.map(lambda a: a.shape, clip(updated)).

What to learn next#

• Random numbers and PRNG keys — the last plumbing before building real models.
• Building models with Flax — modules whose parameters are exactly these trees.
• Optax — tree-shaped optimizer states in action.