Nested Data Structures: Mastering JAX PyTrees

Introduction

Welcome to the lesson 4 of our course "Advanced JAX: Transformations for Speed & Scale"! It's fantastic to continue this journey with you as we explore increasingly sophisticated aspects of JAX. In our previous lessons, you mastered JAX's functional approach to randomness with explicit key management, learned how jax.vmap enables elegant automatic vectorization, and discovered how jax.shard_map allows you to harness parallel computing across multiple devices. These transformations all share a common foundation: they work seamlessly with JAX's flexible data structures. In this lesson, we'll explore JAX PyTrees — arbitrary nested Python containers whose "leaves" are typically JAX arrays. PyTrees are fundamental to JAX's design philosophy, serving as the universal data structure that powers everything from model parameters and optimizer states to complex nested computations. Just as you learned to explicitly manage randomness and efficiently batch operations, understanding PyTrees will give you precise control over how JAX transformations handle structured data. By the end of this lesson, you'll understand how to create and work with PyTree structures, use jax.tree_util.tree_map to apply functions uniformly across all array leaves, and recognize why PyTrees make JAX transformations so powerful and composable. This knowledge will be essential as we continue building toward more complex applications in machine learning and scientific computing.

What Are PyTrees?

Before diving into code, let's establish a clear understanding of what PyTrees are and why they're so central to JAX. A PyTree (short for "Python Tree") is JAX's term for any nested structure of Python containers, such as lists, tuples, dictionaries, or even custom classes you define, whose ultimate "leaves" are typically JAX arrays or other JAX-compatible data types. The beauty of PyTrees lies in their flexibility and ubiquity. Consider these common scenarios where PyTrees are indispensable: Neural network parameters: A dictionary like {'weights': weight_matrix, 'bias': bias_vector} naturally represents a layer's parameters. Optimizer states: Nested structures tracking momentum, gradients, and other variables crucial for model training. Structured datasets: Lists of samples, where each sample might be a dictionary containing multiple fields or measurements. Complex model architectures: Deeply nested dictionaries representing different components of sophisticated models, like encoders and decoders. What makes PyTrees special is that all JAX transformations understand them natively. When you apply jax.jit, jax.grad, jax.vmap, or jax.shard_map to a function that takes PyTree arguments, JAX automatically handles the nested structure. The transformation applies to each array leaf appropriately while preserving the overall container structure. This seamless integration is what makes JAX code both powerful and intuitive. The key insight is this: PyTrees allow you to organize your data naturally using familiar Python structures, while JAX's transformations work transparently with whatever organization you choose. Let's see this in action.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

import jax
import jax.numpy as jnp
import jax.tree_util as jtu # Standard alias for jax.tree_util

# Create a simple PyTree (dictionary representing layer parameters)
key = jax.random.key(202) # PRNG key for reproducibility
key_w, key_b = jax.random.split(key) # Split key for weights and bias

simple_layer_params = {
    'weights': jax.random.normal(key_w, (3, 2)),  # A 3x2 JAX array for weights
    'bias': jax.random.normal(key_b, (2,))        # A JAX array of size 2 for bias
}
print("Simple PyTree (layer parameters):")
print(simple_layer_params)

import jax
import jax.numpy as jnp
import jax.tree_util as jtu # Standard alias for jax.tree_util

# Create a simple PyTree (dictionary representing layer parameters)
key = jax.random.key(202) # PRNG key for reproducibility
key_w, key_b = jax.random.split(key) # Split key for weights and bias

simple_layer_params = {
    'weights': jax.random.normal(key_w, (3, 2)),  # A 3x2 JAX array for weights
    'bias': jax.random.normal(key_b, (2,))        # A JAX array of size 2 for bias
}
print("Simple PyTree (layer parameters):")
print(simple_layer_params)

text

Simple PyTree (layer parameters):
{'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}

Simple PyTree (layer parameters):
{'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}

Python

# Create a more complex PyTree (list of layers)
# Split the key further for more parameters
key, key_w2, key_b2, key_w3, key_b3 = jax.random.split(key, 5) 

model_params = [
    { # First layer's parameters
        'weights': jax.random.normal(key_w2, (4, 3)),
        'bias': jax.random.normal(key_b2, (3,))
    },
    { # Second layer's parameters
        'weights': jax.random.normal(key_w3, (3, 2)),
        'bias': jax.random.normal(key_b3, (2,))
    }
]
print("\nComplex PyTree (list of layer parameters):")
print(model_params)

# Another example: a nested dictionary representing a more complex model
nested_dict_params = {
    'encoder': {
        'gru_cell': {'kernel': jnp.ones((3,3)), 'bias': jnp.zeros(3)}
    },
    'decoder': {
        'dense_layer': simple_layer_params # Reusing the PyTree from earlier!
    }
}
print("\nComplex PyTree (nested dictionary):")
print(nested_dict_params)

# Create a more complex PyTree (list of layers)
# Split the key further for more parameters
key, key_w2, key_b2, key_w3, key_b3 = jax.random.split(key, 5) 

model_params = [
    { # First layer's parameters
        'weights': jax.random.normal(key_w2, (4, 3)),
        'bias': jax.random.normal(key_b2, (3,))
    },
    { # Second layer's parameters
        'weights': jax.random.normal(key_w3, (3, 2)),
        'bias': jax.random.normal(key_b3, (2,))
    }
]
print("\nComplex PyTree (list of layer parameters):")
print(model_params)

# Another example: a nested dictionary representing a more complex model
nested_dict_params = {
    'encoder': {
        'gru_cell': {'kernel': jnp.ones((3,3)), 'bias': jnp.zeros(3)}
    },
    'decoder': {
        'dense_layer': simple_layer_params # Reusing the PyTree from earlier!
    }
}
print("\nComplex PyTree (nested dictionary):")
print(nested_dict_params)

text

Complex PyTree (list of layer parameters):
[{'weights': Array([[-1.2210863 ,  1.1554664 , -0.10118528],
       [ 0.09493745,  0.82033074, -2.208173  ],
       [ 0.34835505,  0.94139576,  0.3934937 ],
       [-0.79533136,  0.5219249 ,  1.3720461 ]], dtype=float32), 'bias': Array([ 2.0215135, -0.8006127,  0.573482 ], dtype=float32)}, {'weights': Array([[ 0.04348394, -0.8883748 ],
       [-1.2707124 ,  0.3398359 ],
       [-0.20632394,  1.2779288 ]], dtype=float32), 'bias': Array([-1.0617399, -0.8719812], dtype=float32)}]

Complex PyTree (nested dictionary):
{'encoder': {'gru_cell': {'kernel': Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32), 'bias': Array([0., 0., 0.], dtype=float32)}}, 'decoder': {'dense_layer': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}}}

Complex PyTree (list of layer parameters):
[{'weights': Array([[-1.2210863 ,  1.1554664 , -0.10118528],
       [ 0.09493745,  0.82033074, -2.208173  ],
       [ 0.34835505,  0.94139576,  0.3934937 ],
       [-0.79533136,  0.5219249 ,  1.3720461 ]], dtype=float32), 'bias': Array([ 2.0215135, -0.8006127,  0.573482 ], dtype=float32)}, {'weights': Array([[ 0.04348394, -0.8883748 ],
       [-1.2707124 ,  0.3398359 ],
       [-0.20632394,  1.2779288 ]], dtype=float32), 'bias': Array([-1.0617399, -0.8719812], dtype=float32)}]

Complex PyTree (nested dictionary):
{'encoder': {'gru_cell': {'kernel': Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32), 'bias': Array([0., 0., 0.], dtype=float32)}}, 'decoder': {'dense_layer': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}}}

Python

# Define simple functions to apply to leaves
def scale_leaf(leaf_array):
    # We only want to scale JAX arrays, not other potential leaf types
    if isinstance(leaf_array, (jnp.ndarray, jax.Array)):
        return leaf_array * 2.0
    else:
        return leaf_array  # Return non-array leaves unchanged

def add_one_to_leaf(leaf_array):
    # This function assumes all leaves it operates on are JAX arrays
    return leaf_array + 1.0

# Use jax.tree_util.tree_map (aliased as jtu)
print("\nApplying `scale_leaf` (multiply by 2.0) to simple_layer_params:")
scaled_simple_params = jtu.tree_map(scale_leaf, simple_layer_params)
print(scaled_simple_params)

print("\nApplying `add_one_to_leaf` to model_params:")
# model_params only contains JAX arrays as leaves, so add_one_to_leaf is safe
incremented_model_params = jtu.tree_map(add_one_to_leaf, model_params)
print(incremented_model_params)

# Define simple functions to apply to leaves
def scale_leaf(leaf_array):
    # We only want to scale JAX arrays, not other potential leaf types
    if isinstance(leaf_array, (jnp.ndarray, jax.Array)):
        return leaf_array * 2.0
    else:
        return leaf_array  # Return non-array leaves unchanged

def add_one_to_leaf(leaf_array):
    # This function assumes all leaves it operates on are JAX arrays
    return leaf_array + 1.0

# Use jax.tree_util.tree_map (aliased as jtu)
print("\nApplying `scale_leaf` (multiply by 2.0) to simple_layer_params:")
scaled_simple_params = jtu.tree_map(scale_leaf, simple_layer_params)
print(scaled_simple_params)

print("\nApplying `add_one_to_leaf` to model_params:")
# model_params only contains JAX arrays as leaves, so add_one_to_leaf is safe
incremented_model_params = jtu.tree_map(add_one_to_leaf, model_params)
print(incremented_model_params)

text

Applying `scale_leaf` (multiply by 2.0) to simple_layer_params:
{'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}

Applying `add_one_to_leaf` to model_params:
[{'bias': Array([3.0215135 , 0.19938731, 1.573482  ], dtype=float32), 'weights': Array([[-0.22108626,  2.1554666 ,  0.89881474],
       [ 1.0949374 ,  1.8203307 , -1.208173  ],
       [ 1.348355  ,  1.9413958 ,  1.3934937 ],
       [ 0.20466864,  1.521925  ,  2.372046  ]], dtype=float32)}, {'bias': Array([-0.06173992,  0.1280188 ], dtype=float32), 'weights': Array([[ 1.043484  ,  0.11162519],
       [-0.27071238,  1.3398359 ],
       [ 0.7936761 ,  2.2779288 ]], dtype=float32)}]

Applying `scale_leaf` (multiply by 2.0) to simple_layer_params:
{'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}

Applying `add_one_to_leaf` to model_params:
[{'bias': Array([3.0215135 , 0.19938731, 1.573482  ], dtype=float32), 'weights': Array([[-0.22108626,  2.1554666 ,  0.89881474],
       [ 1.0949374 ,  1.8203307 , -1.208173  ],
       [ 1.348355  ,  1.9413958 ,  1.3934937 ],
       [ 0.20466864,  1.521925  ,  2.372046  ]], dtype=float32)}, {'bias': Array([-0.06173992,  0.1280188 ], dtype=float32), 'weights': Array([[ 1.043484  ,  0.11162519],
       [-0.27071238,  1.3398359 ],
       [ 0.7936761 ,  2.2779288 ]], dtype=float32)}]

Python

# tree_map can also take multiple PyTrees if they have the same structure
print("\nApplying element-wise addition to two identical PyTrees:")
params_A = {'w': jnp.array([1., 2.]), 'b': jnp.array([0.1])}
params_B = {'w': jnp.array([3., 4.]), 'b': jnp.array([0.2])}

# The lambda function will receive corresponding leaves from params_A and params_B
summed_params = jtu.tree_map(lambda x, y: x + y, params_A, params_B)
print(f"Params A: {params_A}")
print(f"Params B: {params_B}")
print(f"Summed Params: {summed_params}")

# tree_map can also take multiple PyTrees if they have the same structure
print("\nApplying element-wise addition to two identical PyTrees:")
params_A = {'w': jnp.array([1., 2.]), 'b': jnp.array([0.1])}
params_B = {'w': jnp.array([3., 4.]), 'b': jnp.array([0.2])}

# The lambda function will receive corresponding leaves from params_A and params_B
summed_params = jtu.tree_map(lambda x, y: x + y, params_A, params_B)
print(f"Params A: {params_A}")
print(f"Params B: {params_B}")
print(f"Summed Params: {summed_params}")

Python

# PyTrees can contain non-JAX array leaves (e.g., strings, integers)
pytree_with_non_array = {
    'params': simple_layer_params, # A nested PyTree
    'name': 'my_layer',            # A string leaf
    'id': 123                      # An integer leaf
}
print("\nPyTree with non-array leaves:")
print(pytree_with_non_array)

# Apply scale_leaf. Recall scale_leaf checks isinstance and only scales arrays.
scaled_pytree_mixed = jtu.tree_map(scale_leaf, pytree_with_non_array)
print("Scaled PyTree (non-arrays are preserved):")
print(scaled_pytree_mixed)

# PyTrees can contain non-JAX array leaves (e.g., strings, integers)
pytree_with_non_array = {
    'params': simple_layer_params, # A nested PyTree
    'name': 'my_layer',            # A string leaf
    'id': 123                      # An integer leaf
}
print("\nPyTree with non-array leaves:")
print(pytree_with_non_array)

# Apply scale_leaf. Recall scale_leaf checks isinstance and only scales arrays.
scaled_pytree_mixed = jtu.tree_map(scale_leaf, pytree_with_non_array)
print("Scaled PyTree (non-arrays are preserved):")
print(scaled_pytree_mixed)

text

Applying element-wise addition to two identical PyTrees:
Params A: {'w': Array([1., 2.], dtype=float32), 'b': Array([0.1], dtype=float32)}
Params B: {'w': Array([3., 4.], dtype=float32), 'b': Array([0.2], dtype=float32)}
Summed Params: {'b': Array([0.3], dtype=float32), 'w': Array([4., 6.], dtype=float32)}

PyTree with non-array leaves:
{'params': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}, 'name': 'my_layer', 'id': 123}
Scaled PyTree (non-arrays are preserved):
{'id': 123, 'name': 'my_layer', 'params': {'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}}

Applying element-wise addition to two identical PyTrees:
Params A: {'w': Array([1., 2.], dtype=float32), 'b': Array([0.1], dtype=float32)}
Params B: {'w': Array([3., 4.], dtype=float32), 'b': Array([0.2], dtype=float32)}
Summed Params: {'b': Array([0.3], dtype=float32), 'w': Array([4., 6.], dtype=float32)}

PyTree with non-array leaves:
{'params': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}, 'name': 'my_layer', 'id': 123}
Scaled PyTree (non-arrays are preserved):
{'id': 123, 'name': 'my_layer', 'params': {'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}}

Python

import jax
import jax.numpy as jnp
import jax.tree_util as jtu # Standard alias for jax.tree_util

# Create a simple PyTree (dictionary representing layer parameters)
key = jax.random.key(202) # PRNG key for reproducibility
key_w, key_b = jax.random.split(key) # Split key for weights and bias

simple_layer_params = {
    'weights': jax.random.normal(key_w, (3, 2)),  # A 3x2 JAX array for weights
    'bias': jax.random.normal(key_b, (2,))        # A JAX array of size 2 for bias
}
print("Simple PyTree (layer parameters):")
print(simple_layer_params)

text

Simple PyTree (layer parameters):
{'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}

Python

# Create a more complex PyTree (list of layers)
# Split the key further for more parameters
key, key_w2, key_b2, key_w3, key_b3 = jax.random.split(key, 5) 

model_params = [
    { # First layer's parameters
        'weights': jax.random.normal(key_w2, (4, 3)),
        'bias': jax.random.normal(key_b2, (3,))
    },
    { # Second layer's parameters
        'weights': jax.random.normal(key_w3, (3, 2)),
        'bias': jax.random.normal(key_b3, (2,))
    }
]
print("\nComplex PyTree (list of layer parameters):")
print(model_params)

# Another example: a nested dictionary representing a more complex model
nested_dict_params = {
    'encoder': {
        'gru_cell': {'kernel': jnp.ones((3,3)), 'bias': jnp.zeros(3)}
    },
    'decoder': {
        'dense_layer': simple_layer_params # Reusing the PyTree from earlier!
    }
}
print("\nComplex PyTree (nested dictionary):")
print(nested_dict_params)

text

Complex PyTree (list of layer parameters):
[{'weights': Array([[-1.2210863 ,  1.1554664 , -0.10118528],
       [ 0.09493745,  0.82033074, -2.208173  ],
       [ 0.34835505,  0.94139576,  0.3934937 ],
       [-0.79533136,  0.5219249 ,  1.3720461 ]], dtype=float32), 'bias': Array([ 2.0215135, -0.8006127,  0.573482 ], dtype=float32)}, {'weights': Array([[ 0.04348394, -0.8883748 ],
       [-1.2707124 ,  0.3398359 ],
       [-0.20632394,  1.2779288 ]], dtype=float32), 'bias': Array([-1.0617399, -0.8719812], dtype=float32)}]

Complex PyTree (nested dictionary):
{'encoder': {'gru_cell': {'kernel': Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32), 'bias': Array([0., 0., 0.], dtype=float32)}}, 'decoder': {'dense_layer': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}}}

Python

# Define simple functions to apply to leaves
def scale_leaf(leaf_array):
    # We only want to scale JAX arrays, not other potential leaf types
    if isinstance(leaf_array, (jnp.ndarray, jax.Array)):
        return leaf_array * 2.0
    else:
        return leaf_array  # Return non-array leaves unchanged

def add_one_to_leaf(leaf_array):
    # This function assumes all leaves it operates on are JAX arrays
    return leaf_array + 1.0

# Use jax.tree_util.tree_map (aliased as jtu)
print("\nApplying `scale_leaf` (multiply by 2.0) to simple_layer_params:")
scaled_simple_params = jtu.tree_map(scale_leaf, simple_layer_params)
print(scaled_simple_params)

print("\nApplying `add_one_to_leaf` to model_params:")
# model_params only contains JAX arrays as leaves, so add_one_to_leaf is safe
incremented_model_params = jtu.tree_map(add_one_to_leaf, model_params)
print(incremented_model_params)

text

Applying `scale_leaf` (multiply by 2.0) to simple_layer_params:
{'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}

Applying `add_one_to_leaf` to model_params:
[{'bias': Array([3.0215135 , 0.19938731, 1.573482  ], dtype=float32), 'weights': Array([[-0.22108626,  2.1554666 ,  0.89881474],
       [ 1.0949374 ,  1.8203307 , -1.208173  ],
       [ 1.348355  ,  1.9413958 ,  1.3934937 ],
       [ 0.20466864,  1.521925  ,  2.372046  ]], dtype=float32)}, {'bias': Array([-0.06173992,  0.1280188 ], dtype=float32), 'weights': Array([[ 1.043484  ,  0.11162519],
       [-0.27071238,  1.3398359 ],
       [ 0.7936761 ,  2.2779288 ]], dtype=float32)}]

Python

# tree_map can also take multiple PyTrees if they have the same structure
print("\nApplying element-wise addition to two identical PyTrees:")
params_A = {'w': jnp.array([1., 2.]), 'b': jnp.array([0.1])}
params_B = {'w': jnp.array([3., 4.]), 'b': jnp.array([0.2])}

# The lambda function will receive corresponding leaves from params_A and params_B
summed_params = jtu.tree_map(lambda x, y: x + y, params_A, params_B)
print(f"Params A: {params_A}")
print(f"Params B: {params_B}")
print(f"Summed Params: {summed_params}")

Python

# PyTrees can contain non-JAX array leaves (e.g., strings, integers)
pytree_with_non_array = {
    'params': simple_layer_params, # A nested PyTree
    'name': 'my_layer',            # A string leaf
    'id': 123                      # An integer leaf
}
print("\nPyTree with non-array leaves:")
print(pytree_with_non_array)

# Apply scale_leaf. Recall scale_leaf checks isinstance and only scales arrays.
scaled_pytree_mixed = jtu.tree_map(scale_leaf, pytree_with_non_array)
print("Scaled PyTree (non-arrays are preserved):")
print(scaled_pytree_mixed)

text

Applying element-wise addition to two identical PyTrees:
Params A: {'w': Array([1., 2.], dtype=float32), 'b': Array([0.1], dtype=float32)}
Params B: {'w': Array([3., 4.], dtype=float32), 'b': Array([0.2], dtype=float32)}
Summed Params: {'b': Array([0.3], dtype=float32), 'w': Array([4., 6.], dtype=float32)}

PyTree with non-array leaves:
{'params': {'weights': Array([[-0.53820986, -2.3059752 ],
       [ 1.0497485 , -0.30329454],
       [-2.535212  , -0.39521727]], dtype=float32), 'bias': Array([-1.2210863,  1.1554664], dtype=float32)}, 'name': 'my_layer', 'id': 123}
Scaled PyTree (non-arrays are preserved):
{'id': 123, 'name': 'my_layer', 'params': {'bias': Array([-2.4421725,  2.3109329], dtype=float32), 'weights': Array([[-1.0764197 , -4.6119504 ],
       [ 2.099497  , -0.6065891 ],
       [-5.070424  , -0.79043454]], dtype=float32)}}