Skip to content

← JAX Foundations step 16 of 25

Easy Primitives

Pytree Leaves

Jump to the task

Why this matters

A pytree is a value viewed as nested containers and leaves. JAX knows how to traverse ordinary dictionaries, lists and tuples. Arrays are leaves: a matrix is one leaf, not a subtree of its elements.

This lets a model keep parameters in a readable structure while utilities visit the arrays inside it. Tree utilities can also handle non-array leaves, though not every leaf type is a valid JIT input.

Worked mini-example

import jax
import jax.numpy as jnp

params = {
    "w": jnp.array([[1.0, 2.0], [3.0, 4.0]]),
    "b": jnp.array([5.0, 6.0]),
}
leaves = jax.tree_util.tree_leaves(params)
# Ordinary dict keys are sorted: leaves are [params["b"], params["w"]].
total = sum(jnp.sum(leaf) for leaf in leaves)  # 11 + 10 = 21

The result of tree_leaves is a Python list containing two arrays. Their shapes remain (2,) and (2, 2). It does not concatenate them into one vector or flatten the dimensions of either array.

Common pitfalls

  • A single array is a valid pytree with one leaf, even if it is empty.
  • By default, None contributes no leaves. An is_leaf rule can change that behavior.
  • An unregistered custom class is normally treated as one leaf. Register a custom node when JAX should traverse its attributes.
  • tree_map visits tree leaves. vmap batches a function over array axes; these are different operations.
  • Optimizer state can contain parameter-shaped arrays plus counters or other data. It need not have exactly the parameter tree's structure.

Problem

Implement sum_pytree_leaves(leaves). Despite the parameter name, this exercise supplies one one-dimensional JAX array, not the list returned by tree_leaves.

Return the sum of all elements as a scalar. For [1, 2, 3], return 6; for an empty array, return zero.

This small exercise checks a reduction. The worked example shows the extra step needed to sum a real nested tree of arrays. The diagram illustrates nesting and leaves, not the input format used by this exercise.

Loading visualization…