Skip to content

← JAX Foundations step 17 of 25

Medium Primitives

Pytree Map

Jump to the task

Why this matters

jax.tree_util.tree_map(fn, tree) calls fn on each leaf. You can scale every parameter array without writing a separate loop for each level of a nested model.

When the function returns an array for each array leaf, the result keeps the same container structure. Returning a list or tuple at a leaf can introduce new nested structure instead.

Worked mini-example

from jax import tree_util as jtu
import jax.numpy as jnp

params = {
    "w": jnp.array([[1.0, 2.0], [3.0, 4.0]]),
    "b": jnp.array([5.0, 6.0]),
}
scaled = jtu.tree_map(lambda leaf: leaf * 2.0, params)
# scaled["w"] = [[2, 4], [6, 8]]
# scaled["b"] = [10, 12]

combined = jtu.tree_map(lambda a, b: a + b, params, scaled)
# combined["w"] = [[3, 6], [9, 12]]
# combined["b"] = [15, 18]

The multiplication and addition operate on whole arrays. Tree traversal chooses which arrays to pass to each call.

Common pitfalls

  • For ordinary pairwise array updates, use matching tree structures. More generally, the first tree may be a prefix of later trees: a leaf in the first can match a whole subtree in a later argument.
  • Leaves need not be arrays. The mapped function must support the leaf values it actually receives.
  • tree_map is a Python utility, not a compilation boundary. Under JIT its traversal happens during tracing; the recorded array operations then run in the compiled computation.
  • None contributes no leaves by default. That fact alone does not freeze a parameter; the optimizer must interpret its mask or update structure consistently.

Problem

Implement scale_all_leaves(leaves_concat, leaf_sizes, scale).

This exercise uses a one-dimensional array of concatenated values. Multiply every element by the scalar scale and return the same shape. leaf_sizes records how many values each original leaf contributed, supplied as integer-valued floats. You do not need those boundaries to apply one scale to every element.

For leaves_concat = [4, -2, 0], leaf_sizes = [1, 2] and scale = 3, return [12, -6, 0]. An empty input produces an empty output.

This packed-array contract is specific to the exercise. It does not test traversal of a nested tree; use the worked example to try that API. Leaf sizes alone also do not record the original container keys or each leaf's dimensions.

Loading visualization…