Nodes you register

Wrap two arrays in a plain dataclass, flatten it, and there is exactly one leaf: the object itself. PyTreeDef(*) is JAX reporting that it found nothing it knew how to walk into, and it reports that quietly, at flatten time, with no error anywhere.

The refusal comes later and from a different direction entirely. jit takes that single leaf, tries to turn it into an abstract array, and fails on the type. The message names the path, which here is just the argument, and suggests static_argnums, which is the wrong fix for a state object holding arrays you want traced.

Chapter 7 puts it as opting in: a class joins the pytree system the moment you register it. This is what the interval before that looks like from the inside, and it is quiet until an actual transformation touches the object.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@dataclass
class State:
    w: jax.Array
    step: int

s = State(jnp.ones(3), 0)
print(jax.tree.structure(s), len(jax.tree.leaves(s)))
try:
    jax.jit(lambda st: st.w * 2)(s)
except TypeError as e:
    print(str(e).split("abstract array. ", 1)[1].split("\n")[0])

# PyTreeDef(*) 1
# The problematic value is of type <class '__main__.State'> and was passed to the function at path st.

register_pytree_node takes a type and two functions. Flatten returns the children plus one more value, and unflatten receives that value back along with the rebuilt children. The reference calls the extra value "some hashable auxiliary data to be stored in the treedef and to be passed to the unflatten_func", and both halves of that phrase carry a requirement.

The printed treedef shows the split in one line. PyTreeDef(CustomNode(Layer[relu], [*, *])) puts the aux value in brackets after the type name and the children in the star positions after it. Everything in the brackets is structure. Everything in the stars is data.

Which side a field lands on decides how it behaves under every transformation. Children get traced, batched and differentiated like any other leaf. Aux data is carried along untouched and handed back to your constructor on the way out, which is why the activation name below survives a round trip through jit without ever becoming a tracer.

import jax
import jax.numpy as jnp

class Layer:
    def __init__(self, w, b, act):
        self.w, self.b, self.act = w, b, act

jax.tree_util.register_pytree_node(
    Layer,
    lambda l: ((l.w, l.b), l.act),        # children first, then the aux slot
    lambda act, children: Layer(*children, act=act),
)

layer = Layer(jnp.ones((2, 2)), jnp.zeros(2), "relu")
print(jax.tree.structure(layer))
print([x.shape for x in jax.tree.leaves(layer)])
doubled = jax.jit(lambda l: Layer(l.w * 2, l.b, l.act))(layer)
print(doubled.act, doubled.w[0, 0])

# PyTreeDef(CustomNode(Layer[relu], [*, *]))
# [(2, 2), (2,)]
# relu 2.0

register_dataclass writes the same split for you from field metadata. Mark a field static and it goes into the aux tuple, so a Cfg with depth=2 prints as Cfg[(2,)] and that 2 is now part of the structure rather than part of the data.

Changing it therefore retraces. The counter below sits at 1 across two calls whose arrays differ, then steps to 2 the moment depth goes from 2 to 3. Chapter 3 and LAB J2 name the components of the jit cache key; this is the pytree component moving, driven by a field you declared static.

That gives you a design rule with a cost attached. A field in the aux slot is free to read inside the function and can be branched on with an ordinary Python if, because it is a real value at trace time. It also multiplies your executables by its number of distinct values, so anything with more than a handful of them belongs in the children.

from dataclasses import dataclass, field

import jax
import jax.numpy as jnp

@jax.tree_util.register_dataclass
@dataclass
class Cfg:
    w: jax.Array
    depth: int = field(metadata=dict(static=True), default=2)

traces = 0

@jax.jit
def scaled(c):
    global traces
    traces += 1
    return c.w * c.depth

print(jax.tree.structure(Cfg(jnp.ones(3), 2)))
scaled(Cfg(jnp.ones(3), 2)); print(traces)
scaled(Cfg(jnp.zeros(3), 2)); print(traces)
scaled(Cfg(jnp.zeros(3), 3)); print(traces)

# PyTreeDef(CustomNode(Cfg[(2,)], [*]))
# 1
# 1
# 2

The API docs state the requirement plainly: "Metadata fields must be static, hashable, immutable objects, as these objects are used to generate JIT cache keys. In particular, metadata fields cannot contain jax.Array or numpy.ndarray objects." Put one there anyway and the first call goes through without complaint.

The second call is where it lands. Looking up the cache means comparing this call's treedef against the stored one, that comparison reaches the aux values, and comparing two numpy arrays produces an array of booleans rather than a yes or no. What surfaces is numpy's ambiguous-truth-value error, raised from inside a lookup you never wrote.

The museum's exhibit on a static argument that cannot be a cache key shows the same requirement arriving through jit's own arguments. This is that requirement reached from a custom node instead, and the fix has the same shape: put a hashable, comparable stand-in in the aux slot, and keep the array itself among the children.

import jax
import jax.numpy as jnp
import numpy as np

class Masked:
    def __init__(self, w, mask):
        self.w, self.mask = w, mask

jax.tree_util.register_pytree_node(
    Masked,
    lambda m: ((m.w,), m.mask),           # a numpy array in the aux slot
    lambda mask, children: Masked(children[0], mask),
)

double = jax.jit(lambda m: m.w * 2)
print(double(Masked(jnp.ones(3), np.array([True, False]))))
try:
    double(Masked(jnp.zeros(3), np.array([True, False])))
except ValueError as e:
    print(e)

# [2. 2. 2.]
# The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Unflatten does not only run on your arrays. JAX rebuilds your node whenever it needs a copy of the structure, and whatever it happens to be carrying goes into the leaf positions. So a constructor that validates its arguments, which is ordinary good Python, is the wrong constructor to point unflatten at.

The class below refuses anything that is not a jax.Array. Under jit it survives, because the leaves are tracers and tracers are arrays. Under vmap it is handed a bare object, and under eval_shape a ShapeDtypeStruct, and it raises in both.

JAX's own source shows why one of those placeholders exists. In jax/_src/tree_util.py, equality_errors_pytreedef builds a sentinel whose repr is the string pytree leaf, unflattens both treedefs with that sentinel repeated once per leaf, and walks the two results to produce a readable diff. Your unflatten function runs during that walk, with the sentinel in hand and no arrays anywhere.

So keep unflatten to assembly. No validation, no jnp.asarray, no shape checks, no normalizing of defaults. Put all of that in a separate factory that your own code calls, and let the registered constructor accept whatever it is given.

import jax
import jax.numpy as jnp

class Strict:
    def __init__(self, w):
        if not isinstance(w, jax.Array):
            raise TypeError(f"Strict wants an array, got {type(w).__name__}")
        self.w = w

jax.tree_util.register_pytree_node(
    Strict, lambda s: ((s.w,), None), lambda aux, children: Strict(*children)
)

s = Strict(jnp.ones(3))
print(jax.jit(lambda x: Strict(x.w * 2))(s).w)          # tracers and arrays: fine
for label, call in [
    ("vmap", lambda: jax.vmap(lambda x: Strict(x.w * 2))(Strict(jnp.ones((4, 3))))),
    ("eval_shape", lambda: jax.eval_shape(lambda x: Strict(x.w * 2), s)),
]:
    try:
        call()
    except TypeError as e:
        print(label, "->", e)

# [2. 2. 2.]
# vmap -> Strict wants an array, got object
# eval_shape -> Strict wants an array, got ShapeDtypeStruct