Skip to content

← JAX Conceptual Deep-Dives step 2 of 15

Medium Primitives

tree_flatten / tree_unflatten Round-Trip

Jump to the task

Why this matters

tree_flatten separates a pytree into a list of leaves and a description of its containers, called a treedef. tree_unflatten places a list of values into that structure again.

This is useful when you want to transform the leaves as a list while keeping the original dictionary keys, tuple positions and other nesting.

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, treedef = jax.tree_util.tree_flatten(params)
# Sorted dict keys give [params["b"], params["w"]].
rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)

doubled = jax.tree_util.tree_unflatten(
    treedef, [leaf * 2 for leaf in leaves]
)
# doubled["b"] = [10, 12]
# doubled["w"] = [[2, 4], [6, 8]]

The matrix remains one leaf. Flattening a tree does not flatten array dimensions or concatenate array elements.

Common pitfalls

  • Preserve leaf order. Reordering values changes which key or position receives each value; it does not rearrange the treedef.
  • Supply the expected number of leaves. Too many or too few causes an error.
  • A treedef is not an array shape or dtype validator. With ordinary containers, replacing a leaf with an array of a different shape can still reconstruct successfully.
  • A different treedef with the same leaf count can also accept the list. Successful reconstruction does not prove that you used the intended structure.
  • Flattening is not itself a file format. Saving a model also requires a way to store the leaves and reconstruct its structure.

Problem

Implement flatten_then_unflatten_sum(packed_inputs) for a one-dimensional JAX array:

  1. Call jax.tree_util.tree_flatten to obtain leaves and a treedef.
  2. Rebuild with jax.tree_util.tree_unflatten(treedef, leaves).
  3. Return the sum of the rebuilt array as a scalar.

This input is a single-leaf pytree, so the leaf list contains just packed_inputs. For [2, 3], return 5. An empty array is still one leaf and its sum is zero.

The numerical tests check the sum, not whether you performed the round-trip. Use both tree operations to practice the API. The diagram shows a more complex tree; it is separate from the single-array input used here.

Loading visualization…