jax.tree.map — Transforming Pytrees

jax.tree.map is the workhorse function for pytrees. It applies a function to every leaf in a pytree, preserving the tree structure. Think of it like Python's built-in map(), but for arbitrarily nested structures.

import jax
import jax.numpy as jnp

params = {
    'layer1': {'w': jnp.ones((3, 4)), 'b': jnp.zeros(4)},
    'layer2': {'w': jnp.ones((4, 2)), 'b': jnp.zeros(2)},
}

# Apply a function to every leaf
scaled = jax.tree.map(lambda x: x * 2.0, params)
print(scaled['layer1']['w'][0, 0])  # 2.0

# Count parameters
total_params = sum(jax.tree.map(lambda x: x.size, params).values()
                   for layer in jax.tree.leaves(
                       jax.tree.map(lambda x: x.size, params)))
# Better way:
total = sum(x.size for x in jax.tree.leaves(params))
print(f"Total parameters: {total}")  # 3*4 + 4 + 4*2 + 2 = 22

Operating on Multiple Pytrees

The real power of jax.tree.map is that it can operate on multiple pytrees simultaneously, matching them leaf-by-leaf:

# SGD update: params = params - lr * grads
def sgd_update(params, grads, lr=0.01):
    return jax.tree.map(lambda p, g: p - lr * g, params, grads)

# Compute gradients (returns a pytree matching params structure)
def loss_fn(params, x, y):
    h = jax.nn.relu(x @ params['layer1']['w'] + params['layer1']['b'])
    pred = h @ params['layer2']['w'] + params['layer2']['b']
    return jnp.mean((pred - y) ** 2)

x = jnp.ones((5, 3))
y = jnp.ones((5, 2))
grads = jax.grad(loss_fn)(params, x, y)

# grads has SAME structure as params
new_params = sgd_update(params, grads)
# new_params also has the same structure!

API note: jax.tree.map is the modern API (JAX 0.4.25+). The older jax.tree_util.tree_map still works and does exactly the same thing. You'll see both in existing code. Use jax.tree.map for new code.

# Modern API (recommended)
result = jax.tree.map(fn, tree)

# Legacy API (still works)
result = jax.tree_util.tree_map(fn, tree)

# Other useful functions in the modern API
leaves = jax.tree.leaves(tree)            # flat list of leaves
structure = jax.tree.structure(tree)      # just the tree structure
flat, treedef = jax.tree.flatten(tree)    # leaves + structure