Pytree Utilities: flatten, unflatten, and Debugging

Sometimes you need to go beyond tree.map and work with the flat list of leaves directly — for example, to concatenate all parameters into a single vector, or to debug which parameter has NaN values.

Flatten and Unflatten

import jax
import jax.numpy as jnp

params = {
    'layer1': {'w': jnp.array([[1.0, 2.0], [3.0, 4.0]]), 'b': jnp.array([0.1, 0.2])},
    'layer2': {'w': jnp.array([[5.0], [6.0]]), 'b': jnp.array([0.3])},
}

# Flatten: get leaves and structure separately
leaves, treedef = jax.tree.flatten(params)
print(f"Number of leaves: {len(leaves)}")  # 4
print(f"Tree structure: {treedef}")

# You can reconstruct the original tree
params_rebuilt = treedef.unflatten(leaves)
# params_rebuilt == params (same structure and values)

# Useful: concatenate all params into one vector
all_params = jnp.concatenate([x.ravel() for x in leaves])
print(f"Total params as one vector: {all_params.shape}")  # (9,)

Debugging with flatten_with_path

When something goes wrong (NaN gradients, unexpected shapes), you need to know which parameter is the problem. jax.tree.flatten_with_path gives you the path to each leaf:

# Get paths alongside leaves
path_leaves, treedef = jax.tree.flatten_with_path(params)

for path, leaf in path_leaves:
    path_str = jax.tree_util.keystr(path)
    print(f"{path_str}: shape={leaf.shape}, dtype={leaf.dtype}")

# Output:
# ['layer1']['b']: shape=(2,), dtype=float32
# ['layer1']['w']: shape=(2, 2), dtype=float32
# ['layer2']['b']: shape=(1,), dtype=float32
# ['layer2']['w']: shape=(2, 1), dtype=float32

# Find NaN parameters
def check_for_nans(params):
    path_leaves, _ = jax.tree.flatten_with_path(params)
    for path, leaf in path_leaves:
        if jnp.any(jnp.isnan(leaf)):
            print(f"NaN found at {jax.tree_util.keystr(path)}!")

# Check parameter statistics
def param_summary(params):
    path_leaves, _ = jax.tree.flatten_with_path(params)
    for path, leaf in path_leaves:
        name = jax.tree_util.keystr(path)
        print(f"{name}: mean={leaf.mean():.4f}, std={leaf.std():.4f}, "
              f"shape={leaf.shape}")

param_summary(params)
# ['layer1']['b']: mean=0.1500, std=0.0500, shape=(2,)
# ['layer1']['w']: mean=2.5000, std=1.1180, shape=(2, 2)
# ...

Here's a comparison with how PyTorch handles the same task:

# PyTorch equivalent: model.state_dict()
# state_dict = model.state_dict()
# for name, param in state_dict.items():
#     print(f"{name}: {param.shape}")
# Output: "layer1.weight: torch.Size([4, 3])"

# JAX equivalent using pytrees:
path_leaves, _ = jax.tree.flatten_with_path(params)
for path, leaf in path_leaves:
    print(f"{jax.tree_util.keystr(path)}: {leaf.shape}")
# Output: "['layer1']['w']: (3, 4)"

# Both give named access to all parameters — different syntax, same idea