What is a pytree in JAX?

A pytree is either a leaf or a registered container whose children are themselves pytrees, and the recursion stops when everything is a leaf. A leaf is anything the registry does not recognise: an array, a Python scalar, a string, an instance of a class nobody registered. The containers enabled out of the box are tuple, list, dict, None and namedtuple, and you can read that list off the constructor calls in tree_util.py, where each alternative registry is built by passing enable_none, enable_tuple, enable_namedtuple, enable_list and enable_dict explicitly. Flattening returns two things: the list of leaves in traversal order, and a PyTreeDef, which is the shape with the values removed. Unflattening is the inverse and takes new leaves, which is what makes the protocol useful: a transform never has to know what your container was, it works on the flat list and hands the shape back at the end. Two details in the traversal surprise people. Dictionary children are visited in sorted key order rather than insertion order, so the leaf order is not the order you typed. And the traversal is post-order internally, which is why a treedef can be compared node by node as a flat list rather than as a tree.

Why is None not a leaf in JAX pytrees?

Because None is registered as a node with zero children, not as a value. Flattening None gives you an empty leaf list and a treedef that still remembers None was there, so unflattening puts it back. That is what lets you write a parameter tree with an optional field and have grad ignore the field rather than trip over it. The file offers an explicit opt-out. Alongside the default registry it builds a second one whose comment reads that it is a copy of the default registry, where None is a leaf, and the only constructor argument that differs is enable_none set to False. Public functions like jax.tree.flatten always use the default registry; the none-leaf one is there for internal paths that genuinely want None as a value. There is a related trap in the query helpers. treedef_is_leaf asks only whether the treedef has a single node, so it answers True for the treedef of None even though None has no leaves at all. The unexported treedef_is_strict_leaf is the one that also requires exactly one leaf, and it is the one internal code uses when it means what you probably mean.

How do I register a custom class as a JAX pytree?

Six public entry points, and they all end in the same place. register_pytree_node is the foundation: you give it the type, a flatten function returning children plus hashable auxiliary data, an unflatten function taking that auxiliary data and the children back, and optionally a key-aware flatten. Its body is four lines, and they are worth knowing: it loops over a module-level tuple of all four registries and registers the node in every one of them, then records a Python-side entry as well. Register through the public API and the type is visible to the ordinary tree functions, the jit fast-dispatch path and the tracing internals simultaneously; call register_node on one registry yourself and the type is a node on one path and a leaf on another, which produces failures that make no sense. The other five are conveniences on top. register_pytree_node_class is a class decorator expecting tree_flatten and tree_unflatten methods. register_pytree_with_keys takes a flatten that also yields a key per child, and synthesises the plain flatten for you if you do not supply one. register_pytree_with_keys_class is its decorator form. register_dataclass takes field names instead of functions. register_static registers a type with zero children and the instance itself as the auxiliary data, which makes it fully static under jit.

What does register_dataclass do that register_pytree_node_class does not?

It skips the Python callbacks. Instead of storing your flatten and unflatten closures, it calls register_dataclass_node on each registry with two plain lists of field names, and the C++ extension reads the fields with getattr directly and rebuilds the object by keyword. On a hot jit path that removes a Python round trip per call per argument. The split between the two lists is the part that matters for correctness: data fields become children, so they are traced; meta fields become the auxiliary data stored in the treedef, so they are static and end up in the jit cache key. The docstring is blunt that metadata fields must be static, hashable, immutable objects and cannot contain arrays. Since the field-inference change you can omit both lists on a dataclass and the function derives them, treating every field as data unless its dataclasses.field metadata carries static set to True, and skipping anything named in drop_fields. Two guards then run. On a dataclass it computes the set of fields with init True, removes the dropped ones, and refuses if your two lists are not exactly that set, naming the missing and the unexpected fields in the message. And it refuses any field that appears in both lists.

Why do JAX key paths show a flat index instead of my field names?

Because the type was registered without a key-aware flatten. The registration API takes an optional fourth argument, a flatten function that yields a key alongside each child, and when it is absent the C++ side falls back to a backwards-compatibility path that hands out a FlattenedIndexKey for each child, numbered by position. That key prints as a flat index in angle brackets rather than a field name, so every error message and every keystr call over your type says position rather than meaning. The fix is to register with register_pytree_with_keys, or its class decorator, and emit a GetAttrKey per attribute; dataclass registration already does the equivalent for you, since the extension knows the field names. Worth knowing what the built-in keys look like while you are there. A sequence index prints as the index in square brackets, a dictionary key as the repr of the key in square brackets, and an attribute as a dot followed by the name, so a path joins into something like a chain of bracketed keys. keystr also takes a simple flag that drops the brackets and quotes for a compact form, and its own docstring warns that the compact form is ambiguous.

How Every JAX Transform Unpacks Your Data: tree_util.py, Line by Line

jit, grad, vmap and shard_map all take whatever nested mess of dicts, lists and custom classes you hand them, and none of them contains a line of code for walking a container. They do not need one. Every transform calls the same pair of functions first: flatten the argument into a flat list of values plus a description of the shape it came out of, transform the flat list, put the shape back. This is the 1385-line file that defines that protocol, read at a pinned commit of jax-ml/jax, with the eleven regions that matter for flattening and registration walked line by line and the rest left in the pane. Four separate registries built at import, not one, and why a type registered through the public API lands in all four at once. What the C++ extension actually does when you call flatten, including the two rules people get wrong: the is_leaf predicate is consulted before the type lookup, on every node including the root, and dictionary children are visited in sorted key order, so the leaf order is not the order you wrote. Six ways to register a class and what each one costs. Why None is a node with zero children rather than a leaf, and what that does to a treedef. Why a treedef that prints identically to another can still compare unequal. The signature exhibit flattens a nested container you edit, computing the leaf list, the key paths and the treedef string the way the extension builds them.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of How Every JAX Transform Unpacks Your Data: tree_util.py, Line by Line — built to be understood, not skimmed.

Questions

What is a pytree in JAX?: A pytree is either a leaf or a registered container whose children are themselves pytrees, and the recursion stops when everything is a leaf. A leaf is anything the registry does not recognise: an array, a Python scalar, a string, an instance of a class nobody registered. The containers enabled out of the box are tuple, list, dict, None and namedtuple, and you can read that list off the constructor calls in tree_util.py, where each alternative registry is built by passing enable_none, enable_tuple, enable_namedtuple, enable_list and enable_dict explicitly. Flattening returns two things: the list of leaves in traversal order, and a PyTreeDef, which is the shape with the values removed. Unflattening is the inverse and takes new leaves, which is what makes the protocol useful: a transform never has to know what your container was, it works on the flat list and hands the shape back at the end. Two details in the traversal surprise people. Dictionary children are visited in sorted key order rather than insertion order, so the leaf order is not the order you typed. And the traversal is post-order internally, which is why a treedef can be compared node by node as a flat list rather than as a tree.
Why is None not a leaf in JAX pytrees?: Because None is registered as a node with zero children, not as a value. Flattening None gives you an empty leaf list and a treedef that still remembers None was there, so unflattening puts it back. That is what lets you write a parameter tree with an optional field and have grad ignore the field rather than trip over it. The file offers an explicit opt-out. Alongside the default registry it builds a second one whose comment reads that it is a copy of the default registry, where None is a leaf, and the only constructor argument that differs is enable_none set to False. Public functions like jax.tree.flatten always use the default registry; the none-leaf one is there for internal paths that genuinely want None as a value. There is a related trap in the query helpers. treedef_is_leaf asks only whether the treedef has a single node, so it answers True for the treedef of None even though None has no leaves at all. The unexported treedef_is_strict_leaf is the one that also requires exactly one leaf, and it is the one internal code uses when it means what you probably mean.
How do I register a custom class as a JAX pytree?: Six public entry points, and they all end in the same place. register_pytree_node is the foundation: you give it the type, a flatten function returning children plus hashable auxiliary data, an unflatten function taking that auxiliary data and the children back, and optionally a key-aware flatten. Its body is four lines, and they are worth knowing: it loops over a module-level tuple of all four registries and registers the node in every one of them, then records a Python-side entry as well. Register through the public API and the type is visible to the ordinary tree functions, the jit fast-dispatch path and the tracing internals simultaneously; call register_node on one registry yourself and the type is a node on one path and a leaf on another, which produces failures that make no sense. The other five are conveniences on top. register_pytree_node_class is a class decorator expecting tree_flatten and tree_unflatten methods. register_pytree_with_keys takes a flatten that also yields a key per child, and synthesises the plain flatten for you if you do not supply one. register_pytree_with_keys_class is its decorator form. register_dataclass takes field names instead of functions. register_static registers a type with zero children and the instance itself as the auxiliary data, which makes it fully static under jit.
What does register_dataclass do that register_pytree_node_class does not?: It skips the Python callbacks. Instead of storing your flatten and unflatten closures, it calls register_dataclass_node on each registry with two plain lists of field names, and the C++ extension reads the fields with getattr directly and rebuilds the object by keyword. On a hot jit path that removes a Python round trip per call per argument. The split between the two lists is the part that matters for correctness: data fields become children, so they are traced; meta fields become the auxiliary data stored in the treedef, so they are static and end up in the jit cache key. The docstring is blunt that metadata fields must be static, hashable, immutable objects and cannot contain arrays. Since the field-inference change you can omit both lists on a dataclass and the function derives them, treating every field as data unless its dataclasses.field metadata carries static set to True, and skipping anything named in drop_fields. Two guards then run. On a dataclass it computes the set of fields with init True, removes the dropped ones, and refuses if your two lists are not exactly that set, naming the missing and the unexpected fields in the message. And it refuses any field that appears in both lists.
Why do JAX key paths show a flat index instead of my field names?: Because the type was registered without a key-aware flatten. The registration API takes an optional fourth argument, a flatten function that yields a key alongside each child, and when it is absent the C++ side falls back to a backwards-compatibility path that hands out a FlattenedIndexKey for each child, numbered by position. That key prints as a flat index in angle brackets rather than a field name, so every error message and every keystr call over your type says position rather than meaning. The fix is to register with register_pytree_with_keys, or its class decorator, and emit a GetAttrKey per attribute; dataclass registration already does the equivalent for you, since the extension knows the field names. Worth knowing what the built-in keys look like while you are there. A sequence index prints as the index in square brackets, a dictionary key as the repr of the key in square brackets, and an attribute as a dot followed by the name, so a path joins into something like a chain of bracketed keys. keystr also takes a simple flag that drops the brackets and quotes for a compact form, and its own docstring warns that the compact form is ambiguous.

How Every JAX Transform Unpacks Your Data: tree_util.py, Line by Line

Questions

Related explainers