A state that earns its shape

Write Adam's two functions out by hand once and the state stops being an object with a name. init returns a count and two trees of zeros built with jax.tree.map over the params, so jax.tree.structure(opt["mu"]) == jax.tree.structure(params) is True by construction, not by convention. update reads that dict and returns a new one.

The leaves for a 512 by 512 weight and a 512 bias come out as ['int32[]', 'float32[512]', 'float32[512, 512]', 'float32[512]', 'float32[512, 512]']: the count, then mu, then nu, with each tree's keys in sorted order. Five leaves for a two-leaf model.

The bytes make the budget concrete. The parameters take 1050624 bytes, the optimizer state takes 2101252, and 2101252 minus twice 1050624 is 4, which is the int32 count and nothing else. So a step holds three copies of the parameter tree before the gradients it computes are counted, which is the memory that donation from lesson two is trying to stop doubling again.

Optax packages this same pair of functions, and chapter 9 names the pieces. The arithmetic above does not change when you switch to it, because the leaves are the same leaves.

import jax
import jax.numpy as jnp

def adam_init(params):
    zeros = jax.tree.map(jnp.zeros_like, params)
    return {"count": jnp.zeros((), jnp.int32), "mu": zeros, "nu": zeros}

def adam_update(grads, state, lr=1e-2, b1=0.9, b2=0.999, eps=1e-8):
    count = state["count"] + 1
    mu = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g, state["mu"], grads)
    nu = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g * g, state["nu"], grads)
    updates = jax.tree.map(
        lambda m, v: -lr * (m / (1 - b1**count)) / (jnp.sqrt(v / (1 - b2**count)) + eps),
        mu,
        nu,
    )
    return updates, {"count": count, "mu": mu, "nu": nu}

params = {"w": jnp.ones((512, 512)), "b": jnp.zeros(512)}
opt = adam_init(params)

print(jax.tree.structure(opt["mu"]) == jax.tree.structure(params))
print([f"{leaf.dtype.name}{list(leaf.shape)}" for leaf in jax.tree.leaves(opt)])

p_bytes = sum(leaf.nbytes for leaf in jax.tree.leaves(params))
o_bytes = sum(leaf.nbytes for leaf in jax.tree.leaves(opt))
print(p_bytes, o_bytes, o_bytes - 2 * p_bytes)

# True
# ['int32[]', 'float32[512]', 'float32[512, 512]', 'float32[512]', 'float32[512, 512]']
# 1050624 2101252 4

Two ways to hold a step number look interchangeable. 0 and jnp.zeros((), jnp.int32) both mean zero, and jax.eval_shape prints the difference in one word: the Python int comes back as ShapeDtypeStruct(shape=(), dtype=int32, weak_type=True), the array without the weak flag.

The usual warning about the Python version is that it recompiles every step. Measured here, it does not. Three calls, two traces, and the second trace arrives with the strong-typed zero, not with a new value: a weak leaf and its own weak output share one executable however many times the number changes.

Donation does not separate them either. Compiled with donate_argnums=0, both forms report 4 alias bytes, so nothing about buffer reuse argues for one over the other.

The reason to write jnp.zeros((), jnp.int32) is the one measurement leaves standing: an aval that stays the same everywhere the state travels. A counter restored from a checkpoint is a strong int32 array, and the first call that mixes it with a weak-typed one changes the key, which is the flip chapter 3 warns about arriving through the state tree rather than through an argument you wrote by hand.

import jax
import jax.numpy as jnp

print(jax.eval_shape(lambda s: s, 0))
print(jax.eval_shape(lambda s: s, jnp.zeros((), jnp.int32)))

traces = []

def advance(step):
    traces.append(1)
    return step + 1

bump = jax.jit(advance)
bump(0)                          # a Python int: weak_type=True
bump(bump(0))                    # its own output, still weak
bump(jnp.zeros((), jnp.int32))   # the same zero, read back from a checkpoint
print(len(traces))

for step in (0, jnp.zeros((), jnp.int32)):
    compiled = jax.jit(advance, donate_argnums=0).lower(step).compile()
    print(compiled.memory_analysis().alias_size_in_bytes, end=' ')
print()

# ShapeDtypeStruct(shape=(), dtype=int32, weak_type=True)
# ShapeDtypeStruct(shape=(), dtype=int32)
# 2
# 4 4

jax.eval_shape(step, state, batch) runs the step abstractly: no compile, no device work, no allocation. What comes back is the state with a ShapeDtypeStruct at every leaf, and jax.tree.structure(out) == jax.tree.structure(s0) is True, which is the cheapest available proof that a step returns the state it was given rather than something the same size.

The leaf listing is in the treedef's own order, ['float32[4]', 'float32[8, 4]', 'float32[4]', 'float32[8, 4]', 'int32[]']: the params dict with its keys sorted, then the optimizer dict, then the step. Any flatten of that treedef produces that order, which is what a restore has to line up against leaf for leaf.

Chapter 12 sets the rule that the unit of checkpointing is the whole state tree, and the PyTorch path measures what each missing piece costs after a restore at /pytorch/the-loop/the-resume-that-matches. What this lesson adds is where the target tree comes from: an abstract state you can build before the first step runs, from the same function that will produce the real one.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

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

def step(state, batch):
    x, y = batch

    def loss_fn(p):
        return jnp.mean((x @ p["w"] + p["b"] - y) ** 2)

    loss, grads = jax.value_and_grad(loss_fn)(state.params)
    mu = jax.tree.map(lambda m, g: 0.9 * m + 0.1 * g, state.opt["mu"], grads)
    params = jax.tree.map(lambda p, m: p - 0.01 * m, state.params, mu)
    return TrainState(params=params, opt={"mu": mu}, step=state.step + 1), loss

p0 = {"w": jnp.ones((8, 4)) * 0.1, "b": jnp.zeros(4)}
s0 = TrainState(params=p0, opt={"mu": jax.tree.map(jnp.zeros_like, p0)}, step=jnp.zeros((), jnp.int32))

out, loss = jax.eval_shape(step, s0, (jnp.ones((16, 8)), jnp.zeros((16, 4))))
print(jax.tree.structure(out) == jax.tree.structure(s0))
print([f"{leaf.dtype.name}{list(leaf.shape)}" for leaf in jax.tree.leaves(out)])
print(loss)

# True
# ['float32[4]', 'float32[8, 4]', 'float32[4]', 'float32[8, 4]', 'int32[]']
# ShapeDtypeStruct(shape=(), dtype=float32)

Put the four kinds of field side by side and the design rule is short enough to hold in your head. A field is either a leaf or part of the structure. Leaves cost memory and can be donated; structure costs executables and cannot.

One row is borrowed rather than proved here. A field declared static lands in the treedef instead of the leaves, and every distinct value it takes buys its own executable, which the pytrees chapter's lesson on registered nodes counts in its section on the aux slot and the cache key. What a train state adds is the last column: a static field is the one kind that cannot be donated, because it never became an array in the first place.

The bottom row is the one that catches people out. A Python scalar in a state tree is a leaf like any other, and it costs nothing extra until the day something hands you the array-typed version of the same number.