A pytree is JAX's name for any nested structure of containers (dicts, lists, tuples, namedtuples) whose leaves are arrays or scalars. If you've ever stored neural network weights in a nested dictionary, you've already been working with pytrees — you just didn't have a name for it.
import jax
import jax.numpy as jnp
# This is a pytree: nested dicts and lists containing arrays
model_params = {
'encoder': {
'weights': jnp.ones((784, 256)),
'bias': jnp.zeros(256),
},
'decoder': {
'weights': jnp.ones((256, 784)),
'bias': jnp.zeros(784),
},
}
# This is also a pytree: a list of tuples
layers = [
(jnp.ones((3, 4)), jnp.zeros(4)),
(jnp.ones((4, 2)), jnp.zeros(2)),
]
# Even a single array is a (trivial) pytree
single = jnp.array([1.0, 2.0, 3.0])
# You can see the leaves of any pytree
leaves = jax.tree.leaves(model_params)
print(f"Number of leaves: {len(leaves)}") # 4 arrays
print(f"Shapes: {[l.shape for l in leaves]}")
# [(784, 256), (256,), (256, 784), (784,)]The word "pytree" comes from "Python tree." JAX functions like grad, jit, and vmap understand pytrees natively. When you call jax.grad(loss_fn)(params), both params and the returned gradients are pytrees with identical structure — every weight array in params gets a matching gradient array.
💡 Why This Matters
Pytrees are the backbone of JAX's composability. A model's parameters, an optimizer's state, a batch of data, the gradients — they're all pytrees. Every JAX transformation preserves pytree structure automatically. This means you can organize your data however you like (nested dicts, lists of tuples, custom classes) and JAX will work with it without any boilerplate conversion code.
By default, JAX recognizes these container types as pytree nodes:
- dict — keys are sorted for consistency
- list
- tuple
- namedtuple
- None — treated as an empty pytree
Everything else (arrays, numbers, strings, custom objects) is treated as a leaf — an opaque value that JAX doesn't look inside.