Pytrees

JAX functions and transformations operate on arrays, but programs pass around richer structures: a neural network’s parameters might live in a dictionary of arrays with meaningful names, a dataset might be a list of dicts, and so on. JAX has built-in support for such nested structures, which it calls pytrees. This page explains the pytree abstraction, the utilities for working with pytrees, and some common gotchas and patterns.

What is a pytree?

A pytree is a recursive structure: either a leaf, or a container of pytrees. Containers can be lists, tuples, and dicts, nested arbitrarily. A pytree leaf is anything that’s not a container, like an array. A pytree node is anything that’s not a leaf.

Here are some example pytrees, using jax.tree.leaves() to extract the flattened leaves from each:

import jax
import jax.numpy as jnp

example_trees = [
    [1, 'a', object()],
    (1, (2, 3), ()),
    [1, {'k1': 2, 'k2': (3, 4)}, 5],
    {'a': 2, 'b': (2, 3)},
    jnp.array([1, 2, 3]),
]

# Print how many leaves the pytrees have.
for pytree in example_trees:
  leaves = jax.tree.leaves(pytree)
  print(f"{repr(pytree):<45} has {len(leaves)} leaves: {leaves}")

[1, 'a', <object object at 0x7ff52c92f8e0>]   has 3 leaves: [1, 'a', <object object at 0x7ff52c92f8e0>]
(1, (2, 3), ())                               has 3 leaves: [1, 2, 3]
[1, {'k1': 2, 'k2': (3, 4)}, 5]               has 5 leaves: [1, 2, 3, 4, 5]
{'a': 2, 'b': (2, 3)}                         has 3 leaves: [2, 2, 3]
Array([1, 2, 3], dtype=int32)                 has 1 leaves: [Array([1, 2, 3], dtype=int32)]

Which types count as containers is determined by the pytree registry, which by default includes lists, tuples, and dicts (plus a few relatives like namedtuple and OrderedDict). Any object whose type is not in the registry is treated as a leaf. The registry can be extended with user-defined classes; that’s how libraries like Flax and Equinox make entire models into pytrees (Custom pytree nodes).

Conceptually, any pytree can be split into two parts: its leaves (the data) and its treedef (the structure). jax.tree.flatten() performs the split, and jax.tree.unflatten() reassembles:

params = {'W': jnp.zeros((2, 3)), 'b': jnp.zeros(3)}

leaves, treedef = jax.tree.flatten(params)
print(leaves)
print(treedef)
print(jax.tree.unflatten(treedef, leaves))

[Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32), Array([0., 0., 0.], dtype=float32)]
PyTreeDef({'W': *, 'b': *})
{'W': Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32), 'b': Array([0., 0., 0.], dtype=float32)}

This flatten/unflatten decomposition is how JAX transformations support pytrees: internally they operate on the flat list of arrays, then reassemble your structure around the results. Pytrees are tree-like, rather than DAG-like or graph-like, in that we handle them assuming referential transparency and that they can’t contain reference cycles.

Common pytree functions

The pytree utilities live in jax.tree (with lower-level versions in jax.tree_util). The one you’ll use most is jax.tree.map(), which works like Python’s map but operates over entire pytrees:

list_of_lists = [
    [1, 2, 3],
    [1, 2],
    [1, 2, 3, 4]
]

jax.tree.map(lambda x: x * 2, list_of_lists)

[[2, 4, 6], [2, 4], [2, 4, 6, 8]]

jax.tree.map() also supports mapping a function over multiple pytrees at once. The structures must match exactly: lists with the same lengths, dicts with the same keys:

another_list_of_lists = list_of_lists
jax.tree.map(lambda x, y: x + y, list_of_lists, another_list_of_lists)

[[2, 4, 6], [2, 4], [2, 4, 6, 8]]

Other useful functions include jax.tree.reduce() for reductions over leaves and jax.tree.structure() for extracting a treedef; see the jax.tree documentation for the full set.

Pytrees and JAX transformations

All JAX transformations accept functions whose inputs and outputs are pytrees of arrays. Differentiating with respect to a dictionary of parameters just works, and the gradient comes back as a dictionary with the same structure:

def loss(params, x):
  pred = jnp.dot(x, params['W']) + params['b']
  return jnp.sum(pred ** 2)

params = {'W': jnp.ones(3), 'b': 0.5}
x = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])

jax.grad(loss)(params, x)

{'W': Array([137., 181., 225.], dtype=float32),
 'b': Array(44., dtype=float32, weak_type=True)}

This is the pattern that makes JAX practical for machine learning: parameters go in a pytree, jax.grad produces a matching pytree of gradients, and jax.tree.map applies the update. Here’s a complete example, training a small multi-layer perceptron:

import numpy as np

def init_mlp_params(layer_widths):
  params = []
  for n_in, n_out in zip(layer_widths[:-1], layer_widths[1:]):
    params.append(
        dict(weights=np.random.normal(size=(n_in, n_out)) * np.sqrt(2/n_in),
             biases=np.ones(shape=(n_out,)))
    )
  return params

params = init_mlp_params([1, 128, 128, 1])

We can use jax.tree.map to check the shapes of what we built:

jax.tree.map(lambda x: x.shape, params)

[{'biases': (128,), 'weights': (1, 128)},
 {'biases': (128,), 'weights': (128, 128)},
 {'biases': (1,), 'weights': (128, 1)}]

Then define the forward pass, the loss, and the update step:

def forward(params, x):
  *hidden, last = params
  for layer in hidden:
    x = jax.nn.relu(x @ layer['weights'] + layer['biases'])
  return x @ last['weights'] + last['biases']

def loss_fn(params, x, y):
  return jnp.mean((forward(params, x) - y) ** 2)

LEARNING_RATE = 0.0001

def update(params, x, y):
  # `grads` is a pytree with the same structure as `params`.
  grads = jax.grad(loss_fn)(params, x, y)
  # The SGD update is one tree.map over the two matching pytrees.
  return jax.tree.map(
      lambda p, g: p - LEARNING_RATE * g, params, grads
  )

x = np.random.normal(size=(128, 1))
y = x ** 2

for _ in range(100):
  params = update(params, x, y)

print(loss_fn(params, x, y))

0.69971436

(In a real training loop you’d wrap update in jax.jit to make it fast; see Just-in-time compilation.)

Transformation parameters can be pytrees too

Some transformation parameters that refer to inputs, like in_axes and out_axes for jax.vmap(), can themselves be pytrees, matched up against the argument structure. For example, with a function whose second argument is a dict:

vmap(f, in_axes=(0, {"k1": 0, "k2": None}))

maps over the leading axis of the first argument and of k1, while broadcasting k2. These parameter pytrees may also be prefixes of the argument structure, in which case a single value applies to the whole subtree:

vmap(f, in_axes=(0, 0))   # equivalent to (0, {"k1": 0, "k2": 0})
vmap(f, in_axes=0)        # equivalent to (0, {"k1": 0, "k2": 0}) as well

The single-leaf spec in_axes=0 is the familiar default: map everything along its leading axis.

Explicit key paths

Each leaf in a pytree has a key path: the sequence of keys you’d follow to reach it from the root. This is useful for debugging and for anything that needs leaf names, like per-parameter logging. The key-path utilities live in jax.tree_util:

import collections

ATuple = collections.namedtuple("ATuple", ('name',))

tree = [1, {'k1': 2, 'k2': (3, 4)}, ATuple('foo')]
flattened, _ = jax.tree_util.tree_flatten_with_path(tree)

for key_path, value in flattened:
  print(f'Value of tree{jax.tree_util.keystr(key_path)}: {value}')

Value of tree[0]: 1
Value of tree[1]['k1']: 2
Value of tree[1]['k2'][0]: 3
Value of tree[1]['k2'][1]: 4
Value of tree[2].name: foo

jax.tree_util.tree_map_with_path() similarly works like jax.tree.map() with the key path passed as an extra argument.

Custom pytree nodes

By default, any type not in the pytree registry is treated as a leaf, even if it’s a container-like class holding arrays inside:

class Special:
  def __init__(self, x, y):
    self.x = x
    self.y = y

jax.tree.leaves([Special(0, 1), Special(2, 4)])

[<__main__.Special at 0x7ff4f6b58440>, <__main__.Special at 0x7ff4f6b416a0>]

The two Special objects themselves are the leaves. So mapping over what you meant to be the contents fails:

jax.tree.map(lambda x: x + 1, [Special(0, 1), Special(2, 4)])

TypeError: unsupported operand type(s) for +: 'Special' and 'int'

To make your own class act as a container, register it with jax.tree_util.register_pytree_node(), supplying a pair of functions: one that flattens an instance into (children, aux_data), and one that unflattens those pieces back into an instance:

from jax.tree_util import register_pytree_node

class RegisteredSpecial(Special):
  def __repr__(self):
    return f"RegisteredSpecial(x={self.x}, y={self.y})"

def special_flatten(v):
  children = (v.x, v.y)  # the dynamic contents, traversed recursively
  aux_data = None        # static metadata, stored in the treedef
  return children, aux_data

def special_unflatten(aux_data, children):
  return RegisteredSpecial(*children)

register_pytree_node(RegisteredSpecial, special_flatten, special_unflatten)

jax.tree.map(lambda x: x + 1, [RegisteredSpecial(0, 1), RegisteredSpecial(2, 4)])

[RegisteredSpecial(x=1, y=2), RegisteredSpecial(x=3, y=5)]

The division of labor matters: children should hold the dynamic values (arrays and sub-pytrees), while aux_data holds any static metadata. Auxiliary data becomes part of the treedef, which JAX compares and hashes (for example, when deciding whether two pytrees have the same structure), so it must support meaningful equality and hashing.

Once registered, your type works with everything pytrees work with, including transformations. Here’s jax.grad differentiating with respect to a RegisteredSpecial input, returning a matching RegisteredSpecial of gradients:

jax.grad(lambda s: s.x ** 2 + s.y)(RegisteredSpecial(3.0, 4.0))

RegisteredSpecial(x=6.0, y=1.0)

Alternatively, you can define appropriate tree_flatten and tree_unflatten methods on your class and decorate it with register_pytree_node_class():

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class RegisteredSpecial2(Special):
  def __repr__(self):
    return f"RegisteredSpecial2(x={self.x}, y={self.y})"

  def tree_flatten(self):
    children = (self.x, self.y)
    aux_data = None
    return (children, aux_data)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*children)

jax.tree.map(lambda x: x + 1,
             [RegisteredSpecial2(0, 1), RegisteredSpecial2(2, 4)])

[RegisteredSpecial2(x=1, y=2), RegisteredSpecial2(x=3, y=5)]

Some standard Python containers come pre-registered. A NamedTuple subclass, for example, works with no registration at all. But every field becomes a child, including ones you may have meant as metadata:

from typing import NamedTuple, Any

class MyOtherContainer(NamedTuple):
  name: str
  a: Any
  b: Any

jax.tree.leaves([MyOtherContainer('Alice', 1, 2),
                 MyOtherContainer('Bob', 4, 5)])

['Alice', 1, 2, 'Bob', 4, 5]

The names 'Alice' and 'Bob' show up as leaves, which becomes a problem as soon as a transformation tries to treat them as array data.

Registering dataclasses

Unlike NamedTuple subclasses, classes decorated with @dataclass are not automatically pytree nodes. But they’re easy to register, with jax.tree_util.register_dataclass(), and it fixes the metadata problem above, too, by letting you say explicitly which fields are data and which are static metadata:

from dataclasses import dataclass
import functools

@functools.partial(jax.tree_util.register_dataclass,
                   data_fields=['a', 'b'],
                   meta_fields=['name'])
@dataclass
class MyDataclassContainer:
  name: str
  a: Any
  b: Any

jax.tree.leaves([
    MyDataclassContainer('apple', 5.3, 1.2),
    MyDataclassContainer('banana', jnp.zeros(4), -1.0),
])

[5.3, 1.2, Array([0., 0., 0., 0.], dtype=float32), -1.0]

The name field doesn’t appear among the leaves: as a meta_field, it’s carried in the treedef, like aux_data above (and so it must be hashable). This distinction matters again with jax.jit, where meta fields are automatically treated as static arguments; see Marking arguments as static. Instances of MyDataclassContainer can be passed into JIT-ed functions, and name will be treated as static:

@jax.jit
def f(x: MyDataclassContainer | MyOtherContainer):
  return x.a + x.b

# Works fine! `mdc.name` is static.
mdc = MyDataclassContainer('mdc', 1, 2)
y = f(mdc)

Contrast this with MyOtherContainer, the NamedTuple subclass. Since the name field is a pytree leaf, JIT expects it to be convertible to jax.Array, and the following raises an error:

moc = MyOtherContainer('moc', 1, 2)
y = f(moc)

TypeError: Error interpreting argument to <function f at 0x7ff4f6ad3c40> as an abstract array. The problematic value is of type <class 'str'> and was passed to the function at path x.name.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.

Initialization with unexpected values

One caution when writing custom pytree nodes: JAX transformations sometimes build instances of your type with placeholder objects standing in for the real contents, so __init__ and your unflatten function should avoid input validation or array conversion. For example:

class MyTree:
  def __init__(self, a):
    self.a = jnp.asarray(a)

register_pytree_node(MyTree, lambda tree: ((tree.a,), None),
    lambda _, args: MyTree(*args))

tree = MyTree(jnp.arange(5.0))

jax.jacobian(lambda x: x)(tree)  # Error: a placeholder is passed to `MyTree`.

ValueError: None is not a valid value for jnp.array

Here the Jacobian of a function mapping a tree to a tree is defined as a tree of trees, and JAX’s internals build that structure by calling the unflattening recipe with placeholder values, which reach MyTree.__init__.

Potential solution 1:

• The __init__ and __new__ methods of custom pytree classes should generally avoid doing any array conversion or other input validation, or else anticipate and handle these special cases. For example:

class MyTree:
  def __init__(self, a):
    if not (type(a) is object or a is None or isinstance(a, MyTree)):
      a = jnp.asarray(a)
    self.a = a

Potential solution 2:

• Structure your custom tree_unflatten function so that it avoids calling __init__. If you choose this route, make sure that your tree_unflatten function stays in sync with __init__ if and when the code is updated. Example:

def tree_unflatten(aux_data, children):
  del aux_data  # Unused in this class.
  obj = object.__new__(MyTree)
  obj.a = children[0]
  return obj

Common pytree gotchas

Mistaking pytree nodes for leaves

Watch out for accidentally treating nodes as leaves. For example, an array’s .shape is a tuple, which is a pytree node, not a leaf:

a_tree = [jnp.zeros((2, 3)), jnp.zeros((3, 4))]

# Try to make another pytree with ones instead of zeros.
shapes = jax.tree.map(lambda x: x.shape, a_tree)
jax.tree.map(jnp.ones, shapes)

[(Array([1., 1.], dtype=float32), Array([1., 1., 1.], dtype=float32)),
 (Array([1., 1., 1.], dtype=float32), Array([1., 1., 1., 1.], dtype=float32))]

Instead of calling jnp.ones on (2, 3), this called it on 2 and 3 separately, because the tuples became part of the tree structure. The fix depends on the goal: avoid the intermediate tree.map, or make the shape a leaf by converting it to an array.

None is an empty node, not a leaf

jax.tree functions treat None as an empty node, i.e. it has no leaves:

jax.tree.leaves([None, None, None])

[]

To treat None values as leaves, use the is_leaf argument:

jax.tree.leaves([None, None, None], is_leaf=lambda x: x is None)

[None, None, None]

Dictionary keys must be sortable

Dictionaries are flattened by sorted key order, so that pytree structure depends only on the set of keys and not insertion order. That means mixing key types with no ordering between them, like int and str, is an error:

jax.tree.map(lambda x: x + 1, {1: 7, "y": 42})

ValueError: Comparator raised exception while sorting pytree dictionary keys.

If you need unordered keys, collections.OrderedDict flattens in insertion order without sorting, or you can register a custom node type.

Common pytree patterns

Transposing a list of trees into a tree of lists

To turn a list of trees into a tree of lists, the idiom is jax.tree.map with a variadic function:

def tree_transpose(list_of_trees):
  """Converts a list of trees of identical structure into a single tree of lists."""
  return jax.tree.map(lambda *xs: list(xs), *list_of_trees)

# Convert a dataset from row-major to column-major.
episode_steps = [dict(t=1, obs=3), dict(t=2, obs=4)]
tree_transpose(episode_steps)

{'obs': [3, 4], 't': [1, 2]}

For more complex transposes, jax.tree.transpose() lets you specify the inner and outer structure explicitly.

Next steps

With arrays, transformations, and pytrees, you can express most pure computations in JAX. Two topics remain: pseudorandom numbers (Pseudorandom numbers) and stateful computations (Stateful computations).