Nested Data Structures: Mastering JAX PyTrees | CodeSignal Learn

Nested Data Structures: Mastering JAX PyTrees | CodeSignal Learn

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.

Sign up

Join the 1M+ learners on CodeSignal

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

Creating Simple PyTrees

Let's start by creating our first PyTree — a simple dictionary representing neural network layer parameters. This is one of the most common PyTree patterns you'll encounter in machine learning. 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) This code creates what JAX recognizes as a PyTree. As you may recall from our first lesson, we're using JAX's functional random number generation with explicit key splitting. The resulting simple_layer_params dictionary contains two JAX arrays as its values. These arrays are the leaves of our PyTree, while the dictionary structure itself is the container or treedef (tree definition). From JAX's perspective, this isn't just a dictionary; it's a structured container of arrays that can be transformed as a unit. When you pass simple_layer_params to any JAX transformation, JAX will automatically understand its structure and apply operations to both the weights and bias arrays appropriately. The output shows our PyTree structure: 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)} Notice how natural this feels: we're just using a regular Python dictionary, but JAX recognizes it as a structured data container. This is the essence of PyTrees: familiar Python data structures that JAX can understand and manipulate intelligently.

Complex PyTree Structures

PyTrees truly shine when you need more complex nested structures. Let's explore two common patterns: lists of parameter dictionaries (representing multiple layers in a network) and deeply nested dictionaries (representing complex model architectures with multiple components). # 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) These examples demonstrate the flexibility of PyTrees. The model_params variable is a list, where each element is a dictionary representing a layer's parameters. This pattern is common when building multi-layer neural networks. The nested_dict_params shows even deeper nesting, representing a model with distinct 'encoder' and 'decoder' components, where the 'decoder' even reuses our previously defined simple_layer_params PyTree. JAX treats both of these structures as valid PyTrees. The leaves are still the JAX arrays (weights, bias, kernel, etc.), but now they're organized in more sophisticated container hierarchies. This is powerful because JAX transformations will traverse these structures and apply operations to every array leaf, regardless of how deeply nested they are, while preserving the overall organization. The output reveals these nested structures: 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)}}} Notice how we can mix and match container types: lists containing dictionaries, dictionaries containing other dictionaries, and even reuse existing PyTrees as components of larger structures. This composability is fundamental to building complex, maintainable JAX programs.

Transforming PyTrees with tree_map

Now we reach the heart of PyTree functionality: jax.tree_util.tree_map. This versatile function allows you to apply any Python callable (like a function or a lambda) uniformly to all the array leaves in a PyTree, or even multiple PyTrees, while preserving the container structure. It's analogous to Python's built-in map function, but supercharged for nested structures. Let's define some simple transformation functions and see tree_map in action: # 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) Here, jtu.tree_map(scale_leaf, simple_layer_params) takes our scale_leaf function and applies it to every leaf (JAX array) in the simple_layer_params PyTree. The result, scaled_simple_params, is a new PyTree with the exact same structure (a dictionary with 'weights' and 'bias' keys), but where every array has been multiplied by 2.0. Similarly, add_one_to_leaf adds 1.0 to every array element in the more complex model_params structure. The key insight is that tree_map preserves structure perfectly. The output PyTree has exactly the same nesting pattern as the input, but with transformed leaf values: 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)}] Every array in both PyTrees has been transformed, but the dictionary keys, list positions, and overall nesting remain identical. This structure preservation is what makes PyTrees so powerful for complex data transformations.

Advanced tree_map Operations: Multiple PyTrees

tree_map becomes even more powerful when working with multiple PyTrees or handling mixed data types within a PyTree. Let's explore two advanced scenarios that frequently arise in practice. First, tree_map can operate on multiple PyTrees simultaneously, provided they have the same tree structure. This enables element-wise operations between corresponding leaves: # 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}") This demonstrates a crucial capability: when you pass multiple PyTrees (here, params_A and params_B) to tree_map, it applies your function (the lambda lambda x, y: x + y) to corresponding leaves from each PyTree. The lambda receives the 'w' arrays from both params_A and params_B as x and y, adds them element-wise, and the result becomes the 'w' entry in summed_params. The same happens for the 'b' entries. This is extremely useful for operations like adding gradients to parameters or calculating differences between model states.

Advanced tree_map Operations: Non-Array Data

Next, let's see how tree_map handles PyTrees containing non-array data. It's common to have metadata or configuration alongside JAX arrays: # 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) Our scale_leaf function from earlier includes a type check: if isinstance(leaf_array, (jnp.ndarray, jax.Array)):. This ensures it only scales JAX arrays and returns other types of leaves (like the string 'my_layer' or the integer 123) unchanged. This allows PyTrees to gracefully contain metadata, configuration parameters, or other non-numeric data alongside your JAX arrays. The output illustrates both capabilities: 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)}} These advanced patterns — multi-PyTree operations and mixed data types — are essential for real-world applications where you need to combine parameter sets, handle metadata alongside arrays, or implement complex mathematical operations across structured data.

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)}}