One tree in, one tree out

Four parameter arrays, two Adam moments for each of them, and a step counter. Thirteen leaves, and that is the entire run: there is nothing else to save, nothing else to move to a device, nothing else a later step reads. Chapter 7 established that every transformation sees a model as leaves on a tree; this is what the tree looks like once an optimizer and a clock are in it.

Two things in the printout are worth slowing down for. The step counter is int32 while everything above it is float32, so a saver that assumes one dtype for the whole tree is already wrong. And the paths come back in the order b1, b2, w1, w2, not the order the dict was written, because dict keys flatten sorted. Match leaves to names by position and a rename of w1 to a1 silently reorders your checkpoint.

The State type here is a NamedTuple, which JAX already knows how to flatten. A dataclass needs register_dataclass first, which the chapter 7 arc covers along with the sorted-key rule this printout is obeying, and the arc's third lesson depends only on the tree being flattenable, not on which of the two you picked.

import jax
import jax.numpy as jnp
from typing import NamedTuple


class State(NamedTuple):
    params: dict
    m: dict            # Adam's first moment
    v: dict            # Adam's second
    step: jnp.ndarray


def init(key):
    k1, k2 = jax.random.split(key)
    params = {
        "w1": jax.random.normal(k1, (8, 32)) * 0.1, "b1": jnp.zeros(32),
        "w2": jax.random.normal(k2, (32, 1)) * 0.1, "b2": jnp.zeros(1),
    }
    return State(params,
                 jax.tree.map(jnp.zeros_like, params),
                 jax.tree.map(jnp.zeros_like, params),
                 jnp.zeros((), jnp.int32))


state = init(jax.random.key(0))
paths, treedef = jax.tree_util.tree_flatten_with_path(state)
print(len(paths), paths[-1][1].dtype)
print([jax.tree_util.keystr(k) for k, _ in paths[:5]])
print(treedef)

# 13 int32
# [".params['b1']", ".params['b2']", ".params['w1']", ".params['w2']", ".m['b1']"]
# PyTreeDef(CustomNode(namedtuple[State], [{'b1': *, 'b2': *, 'w1': *, 'w2': *},
# {'b1': *, 'b2': *, 'w1': *, 'w2': *}, {'b1': *, 'b2': *, 'w1': *, 'w2': *}, *]))

The step below is the chapter's step with the optimizer written out instead of called. value_and_grad returns the loss and the gradient tree together, the moments update leafwise through tree.map, and the bias correction reads the counter that rides in the same tree. Nothing is stored anywhere; the function takes a State and returns a State.

The batch is drawn inside the step from a key the caller derives with fold_in, so the step number alone decides which 32 rows this step sees. Chapter 8 made that choice for readability under scan. Lesson three spends the whole page on what it buys at resume time.

After twenty steps the loss reads 1.720861 and the jit cache holds exactly one entry. One entry for twenty calls is the steady state a training loop is supposed to reach, and _cache_size is the cheapest way to check it. It is a private helper rather than public API, which is fine for a check you run once and delete.

Twenty calls, one executable. That number is the health check.

LR, B1, B2, EPS = 1e-2, 0.9, 0.999, 1e-8


def predict(p, x):
    h = jax.nn.relu(x @ p["w1"] + p["b1"])
    return h @ p["w2"] + p["b2"]


def loss_fn(p, batch):
    x, y = batch
    return jnp.mean((predict(p, x) - y) ** 2)


def step_fn(state, data, key):
    x, y = data
    idx = jax.random.randint(key, (32,), 0, x.shape[0])
    loss, g = jax.value_and_grad(loss_fn)(state.params, (x[idx], y[idx]))
    t = state.step + 1
    m = jax.tree.map(lambda m, g: B1 * m + (1 - B1) * g, state.m, g)
    v = jax.tree.map(lambda v, g: B2 * v + (1 - B2) * g * g, state.v, g)
    params = jax.tree.map(
        lambda p, m, v: p - LR * (m / (1 - B1 ** t)) / (jnp.sqrt(v / (1 - B2 ** t)) + EPS),
        state.params, m, v,
    )
    return State(params, m, v, t), loss


train_step = jax.jit(step_fn)

x = jax.random.normal(jax.random.key(1), (256, 8))
data = (x, jnp.sum(x, axis=1, keepdims=True))
base = jax.random.key(2)

state = init(jax.random.key(0))
for step in range(20):
    state, loss = train_step(state, data, jax.random.fold_in(base, step))

print(int(state.step), f"{loss:.6f}", type(loss).__name__)
print(train_step._cache_size())

# 20 1.720861 ArrayImpl
# 1

Start the same run with step=0 as a plain Python integer instead of a zero-dimensional array and the loop still produces the right numbers. The cache goes from one entry to three.

Three, not two, and the third one is the part worth understanding. The first call sees a Python int leaf and compiles for it. What that call returns is an int32 array whose weak_type is True, because it was built from a weakly typed input, and a weak int32 is a different signature from the strong int32 the first run used. So the second call compiles again, and only from the third call does the cache stop growing.

Chapter 3 owns the cache-key model, its lesson arc prices a weak Python scalar handed in as an argument, and chapter 11 owns the recompile hunt. What is left for here is that a state tree is the easiest place in a training loop to leak a Python scalar into a leaf, and the symptom is small enough to miss: not a compile per step, just two extra compiles and a cache that will not agree with itself.

s = init(jax.random.key(0))._replace(step=0)          # a Python int, not an array
for step in range(4):
    s, _ = train_step(s, data, jax.random.fold_in(base, step))

print(s.step.dtype, s.step.aval.weak_type, state.step.aval.weak_type)
print(train_step._cache_size())

# int32 True False
# 3

Two lessons elsewhere own donation and this section stays off both. What donate_argnums promises for a single array, what becomes of the array you handed over, and the two mismatches that get a donation refused belong to chapter 11's lesson on it. How a state tree's donation is granted one leaf at a time, and which other names in a loop lose their array when it is, belong to chapter 9's lesson on donating a state tree. Read either one before you add the keyword. What is left over is a failure you can only produce while building an optimizer state, so it belongs here.

Make the zeros once, use them for both Adam moments, and the state looks tidier than the version in the first block. It also has two leaves backed by one buffer, and unsafe_buffer_pointer reports that directly.

Nothing goes wrong until the step is donated. Then one execution is asked to hand the same buffer to two donated slots and refuses with Attempt to donate the same buffer twice, naming the flattened argument index and the earlier use it collided with. The message comes from XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → rather than from JAX, which is why it reads like a runtime status rather than a Python exception.

The two separate tree.map calls of the first block give distinct pointers, and the same donated step then runs and returns step 1. The rule worth carrying past this example: leaves of a donated tree have to be distinct buffers, and sharing that is invisible under normal use turns into a hard failure the moment you promise the runtime it can reuse them.

donating = jax.jit(step_fn, donate_argnums=0)


def init_shared(key):
    k1, k2 = jax.random.split(key)
    params = {
        "w1": jax.random.normal(k1, (8, 32)) * 0.1, "b1": jnp.zeros(32),
        "w2": jax.random.normal(k2, (32, 1)) * 0.1, "b2": jnp.zeros(1),
    }
    zeros = jax.tree.map(jnp.zeros_like, params)     # one tree, used twice
    return State(params, zeros, zeros, jnp.zeros((), jnp.int32))


s = init_shared(jax.random.key(0))
print(s.m["w1"].unsafe_buffer_pointer() == s.v["w1"].unsafe_buffer_pointer())
try:
    donating(s, data, jax.random.fold_in(base, 0))
except Exception as err:
    print(err)

fresh = init(jax.random.key(0))                      # two tree.map calls, two buffers
print(fresh.m["w1"].unsafe_buffer_pointer() == fresh.v["w1"].unsafe_buffer_pointer())
new_state, _ = donating(fresh, data, jax.random.fold_in(base, 0))
print(int(new_state.step))

# True
# INVALID_ARGUMENT: Attempt to donate the same buffer twice in Execute()
# (flattened argument 8, replica 0, partition 0, first use: 4). Toy example for
# this bug: `f(donate(a), donate(a))`.
# False
# 1