Flatten happens first

The forward pass below takes two arguments: a params dict nested two levels deep, and a batch of rows. Trace it and the header line of the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → lists four invars, one per array. Neither the dict nor the word dense is anywhere in the recording.

The two printed lines say the same thing at two altitudes. Chapter 2's recording and chapter 3's lowered StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → agree that this callable takes four flat tensors in one fixed order. Flattening runs on the way in, before tracing starts, and the treedef stays behind on the Python side to rebuild whatever comes back.

Read the order off the header and it is not the order anyone typed. a:f32[2] is dense.b, b:f32[3,2] is dense.w, c:f32[2,1] is out.w, and d:f32[4,3] is the batch. Inside each dict the keys sorted themselves; across the two arguments, positional order held.

The compiler never sees a dict. It sees leaves, in an order you can work out beforehand.

import jax
import jax.numpy as jnp

params = {"dense": {"w": jnp.ones((3, 2)), "b": jnp.zeros(2)},
          "out": {"w": jnp.ones((2, 1))}}

def fwd(p, x):
    h = jnp.tanh(x @ p["dense"]["w"] + p["dense"]["b"])
    return h @ p["out"]["w"]

x = jnp.ones((4, 3))
print(str(jax.make_jaxpr(fwd)(params, x)).split("\n")[0])
print(next(l for l in jax.jit(fwd).lower(params, x).as_text().split("\n") if "@main" in l).strip())

# { lambda ; a:f32[2] b:f32[3,2] c:f32[2,1] d:f32[4,3]. let
# func.func public @main(%arg0: tensor<2xf32>, %arg1: tensor<3x2xf32>, %arg2: tensor<2x1xf32>, %arg3: tensor<4x3xf32>) -> (tensor<4x1xf32> {jax.result_info = ""}) {

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → above put dense.b ahead of dense.w though w was typed first, and nothing in the function chose that. Each container type carries its own flatten function in the registry, and the dict's walks the keys in sorted order, so which array arrives as argument one is settled by the keys before tracing starts. The jit lessons under chapter 3 take the same rule the other way, into what it means for a cache entry; here it decides only the order.

OrderedDict does not inherit that behaviour. It is registered as a node whose auxiliary data is its key tuple, so insertion order is part of the structure and the leaves come out in the order they went in. The first two lines below differ only in the container, and the leaf shapes come out swapped.

A namedtuple is a pytree node without anyone registering it, and its treedef never compares equal to the plain tuple holding the same two fields. The printed forms show why: one is PyTreeDef(CustomNode(namedtuple[P], [*, *])) and the other is PyTreeDef((*, *)), so the class name is carried in the structure rather than looked past.

Keys that cannot be compared to each other stop the sort before any of this happens. A dict mixing an int key with a str key raises while flattening, and the message names the sort rather than the transformation you were calling.

import collections

import jax
import jax.numpy as jnp

w, b = jnp.ones((3, 2)), jnp.zeros(2)
P = collections.namedtuple("P", ["w", "b"])

print([l.shape for l in jax.tree.leaves({"w": w, "b": b})])
print([l.shape for l in jax.tree.leaves(collections.OrderedDict([("w", w), ("b", b)]))])
print(jax.tree.structure(P(w, b)) == jax.tree.structure((w, b)))
try:
    jax.tree.leaves({1: w, "a": b})
except ValueError as e:
    print(e)

# [(2,), (3, 2)]
# [(3, 2), (2,)]
# False
# Comparator raised exception while sorting pytree dictionary keys.

A treedef has a printable form, and reading it beats reasoning about it. PyTreeDef({'cfg': None, 'opt': [*, *], 'w': *}) says there are three leaves, sitting at opt[0], opt[1] and w, and that cfg is present as an empty subtree contributing no leaf. Chapter 7 told you None drops out of the leaves list. The printed form adds that it did not drop out of the structure.

The two counts a treedef reports answer different questions. num_leaves is 3 for that tree and num_nodes is 6, because the nodes tally counts the containers and the None alongside the leaves.

Equality looks at the containers and stops there. Two dicts with the same keys give equal treedefs even when one holds arrays and the other holds Python ints, because the leaves left the structure the moment they were flattened out of it and nothing about them is recorded in what remains.

Unflattening is arity-checked, so a treedef will not silently accept a short list of leaves. Hand jax.tree.unflatten two leaves for a three-leaf structure and it counts them back at you.

import jax
import jax.numpy as jnp

t = {"w": jnp.ones((2, 2)), "cfg": None, "opt": [jnp.zeros(1), 3]}
leaves, treedef = jax.tree.flatten(t)
print(treedef)
print(treedef.num_leaves, treedef.num_nodes)
print(jax.tree.structure({"a": 1, "b": 2}) == jax.tree.structure({"a": jnp.ones(3), "b": jnp.zeros((2, 2))}))
try:
    jax.tree.unflatten(treedef, leaves[:2])
except ValueError as e:
    print(e)

# PyTreeDef({'cfg': None, 'opt': [*, *], 'w': *})
# 3 6
# True
# Too few leaves for PyTreeDef; expected 3, got 2

jax.tree_util.tree_flatten_with_path hands back each leaf together with the route taken to reach it, and keystr prints that route as ['dense']['w']. The same notation shows up in error messages, which is the reason to learn to read it.

Put a string in a params dict and jit refuses, and the refusal locates the offender by path rather than by leaf number. at path p['dense']['act'] finds one bad entry in a tree of hundreds with no printing on your side.

The path survives further down than you might expect. A jitted function returning a nested dict lowers to a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → entry point whose two results carry jax.result_info = "['aux']['norm']" and jax.result_info = "['loss']", recording which slot of the output tree each result belongs to. The tree is gone from the computation and the labels for rebuilding it are not.

import jax
import jax.numpy as jnp

params = {"dense": {"w": jnp.ones((3, 2)), "act": "relu"}}
for path, leaf in jax.tree_util.tree_flatten_with_path(params)[0]:
    print(jax.tree_util.keystr(path), type(leaf).__name__)
try:
    jax.jit(lambda p: p["dense"]["w"] * 2)(params)
except TypeError as e:
    print(str(e).split("abstract array. ", 1)[1].split("\n")[0])

def split(x):
    return {"loss": x.sum(), "aux": {"norm": (x * x).sum()}}

text = jax.jit(split).lower(jnp.ones(4)).as_text()
print(next(l for l in text.split("\n") if "@main" in l).strip())

# ['dense']['act'] str
# ['dense']['w'] ArrayImpl
# The problematic value is of type <class 'str'> and was passed to the function at path p['dense']['act'].
# func.func public @main(%arg0: tensor<4xf32>) -> (tensor<f32> {jax.result_info = "['aux']['norm']"}, tensor<f32> {jax.result_info = "['loss']"}) {