Structure, then leaves

Run a scan whose body returns a carry with one extra key and the message says pytree structure. Run one whose carry starts as int32 and comes back float32 and the message says equal types. Same primitive, same call shape, two entirely different sentences.

The first check compares treedefs and nothing else. The second compares the avals of the leaves, shape and dtype together, and it only runs once the structures already line up. So a structure error means you assembled the wrong containers, and a type error means the containers were right and something inside one of them changed.

Both messages name the component that differs, and the structure one goes as far as printing the symmetric difference of the key sets, {'n'}, so the diffing is done for you. Chapter 6 covers what scan is for. This is what it says when the carry contract breaks, and the wording alone tells you which half broke.

Structure first, then leaves. The message says which pass raised it.

Not every tree complaint is one of these two. vmap can refuse an in_axes spec before either check runs, on the shape of the spec rather than on anything in your arguments, and the vmap lessons under chapter 5 keep that refusal with the others the transformation owns.

import jax
import jax.numpy as jnp

def grows(carry, x):
    return {"m": carry["m"] + x, "n": carry["m"]}, x

def drifts(carry, x):
    return carry + x, x

for body, init in [(grows, {"m": 0.0}), (drifts, jnp.zeros((), jnp.int32))]:
    try:
        jax.lax.scan(body, init, jnp.arange(3.0))
    except TypeError as e:
        print(e)
        print("=" * 8)

# scan body function carry input and carry output must have the same pytree structure, but they differ:
#
# The input carry carry is a <class 'dict'> with 1 child but the corresponding component of the carry output is a <class 'dict'> with 2 children, so the numbers of children do not match, with the symmetric difference of key sets: {'n'}.
#
# Revise the function so that the carry output has the same pytree structure as the carry input.
# ========
# scan body function carry input and carry output must have equal types (e.g. shapes and dtypes of arrays), but they differ:
#
# The input carry carry has type int32[] but the corresponding output carry component has type float32[], so the dtypes do not match.
#
# Revise the function so that all output types (e.g. shapes and dtypes) match the corresponding input types.
# ========

jax.tree.map is not symmetric in its arguments. The first tree is the template, and every later tree is flattened only as far down as that template goes. A deeper structure in argument two is therefore not an error at all: the subtree arrives at your function whole, as one value.

The first printed line shows it happening. The template has a leaf at a, the second tree has a dict there, and the function receives that dict. Nothing checked below the template's depth, because nothing was asked to.

This is why a missing gradient does not raise a structure error. Write None where a layer has no gradient and the template still has a leaf in that position, so None is handed straight to your function and the failure surfaces as arithmetic on a NoneType, several frames away from the tree that caused it.

is_leaf takes the decision back. Pass is_leaf=lambda x: x is None and None stops being an empty subtree for that one call, so your function receives it and decides what a missing gradient means. The structure that comes back out is the template's, unchanged.

import jax
import jax.numpy as jnp

params = {"w": jnp.ones(3), "b": jnp.zeros(3)}
grads = {"w": jnp.ones(3), "b": None}          # a layer with no gradient

print(jax.tree.map(lambda a, b: (a, type(b).__name__), {"a": 1.0}, {"a": {"deeper": 2.0}}))
try:
    jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
except TypeError as e:
    print(e)
out = jax.tree.map(lambda p, g: p if g is None else p - 0.1 * g, params, grads,
                   is_leaf=lambda x: x is None)
print(jax.tree.structure(out))

# {'a': (1.0, 'dict')}
# unsupported operand type(s) for *: 'float' and 'NoneType'
# PyTreeDef({'b': *, 'w': *})

Chapter 4 sets up vjp as the reverse direction and chapter 7 states half of its tree contract, that the gradient comes back carrying the input's treedef. The other half sits on the argument you pass to the pullback, and it has to carry the output's treedef instead.

Return a dict of two things from the function and the pullback takes a dict of two cotangents. Drop one key and the message prints both treedefs side by side, so the diff is done for you.

That symmetry matters as soon as a model returns more than a loss. Whatever shape the forward pass hands back is the shape the backward pass expects to be seeded with, key for key, and a tuple will not stand in for a dict holding the same two entries.

import jax
import jax.numpy as jnp

def f(p):
    return {"y": p["w"] * 2, "z": p["w"].sum()}

out, pullback = jax.vjp(f, {"w": jnp.ones(3)})
print(jax.tree.structure(out))
print(pullback({"y": jnp.ones(3), "z": jnp.ones(())})[0])
try:
    pullback({"y": jnp.ones(3)})
except ValueError as e:
    print(e)

# PyTreeDef({'y': *, 'z': *})
# {'w': Array([3., 3., 3.], dtype=float32)}
# unexpected tree structure of argument to vjp function: got PyTreeDef({'y': *}), but expected to match PyTreeDef({'y': *, 'z': *})