Welcome back everyone, in our previous article, we covered the foundational philosophy of JAX from functional purity to array immutability to basic transformations like jax.grad, jax.vmap and jax.jit. We firsthand saw how treating functions as pure mathematical transformations allows JAX to inspect, trace and optimize our code for hardware like TPUs and GPUs.
However, in the real world, ML architectures are rarely as simple as single arrays passed into scalar functions. Deep neural networks consist of nested layers, distinct weights and biases, optimizer states and hyperparameter dictionaries.
In this article, we shall move beyond flat arrays and explore how JAX uses Pytrees to represent nested data structures without abandoning functional purity or breaking functional rules. By the end of this article, we shall be able to perform functional parameter updates, stack transformations and leveraging object oriented abstractions.
What are Pytrees and Why do they matter
In JAX, a Pytree is a container-like structure built from standard Python collections that contains leaf elements at its tips
A PyTree typically consists of two main components i.e.
Nodes(Branches): These are the structural containers e.g. dicts, lists, tuples, namedtuples or custom registered classes
Leaves: This is the actual data payload located at the innermost points of the tree structure. In JAX, these are typically jax.Array objects, NumPy arrays or scalar values.
An example of a PyTree
# Params of a 2-layer network
params = {
'layer_1': {
'w': jnp.ones((4, 8)),
'b': jnp.zeros((8,))
},
'layer_2': {
'w': jnp.ones((8, 1)),
'b': jnp.zeros((1,))
},
}
Why PyTrees Matter
In object-oriented frameworks like Pytorch and Tensorflow, model parameters are stored inside stateful class attributes like nn.Module . Gradients are calculated by mutating these internal tensors in place using standard attribute lookup.
But because JAX functions must remain pure, parameters must be seperated from execution logic. PyTrees allow us to bundle arbitrarily complex model parameters into a single structured Python object that can be directly passed into transformed functions.
When you pass a PyTree into functions like jax.grad or jax.jit, JAX automatically traverses the tree structure, performs the required transformations on each individual leaf array and returns a new PyTree with the same container structure as the input.
Basic PyTree Operations
Working with nested structures using standard Python for loops can easily turn verbose and inefficient. JAX provides the jax.tree module and jax.tree_util to interact with PyTree functionality.
Inspecting PyTree Structures and Leaves
Before manipulating a PyTree, it's important to inspect its composition. This can be done by either extracting all leaves into a flat list using jax.tree.leaves or inspecting its container skeleton using jax.tree.structure.
import jax
import jax.numpy as jnp
# A simple param PyTree
params = {
'w1': jnp.array([1.0, 2.0], [3.0, 4.0]),
'b1': jnp.array([0.5, -0.5]),
'extra': ('learning_rate', 0.01) #Any primitive types are valid tree elements
}
# Extracting Leaves
leaves = jax.tree.leaves(params)
print("Leaves in params: ")
for leaf in leaves:
print("-", leaf)
# Inspecting the structure
structure = jax.tree.structure(params)
print("\nTree Structure: ")
print(structure)
Output:
Leaves in params:
- [ 0.5 -0.5]
- learning_rate
- 0.01
- [[1. 2.]
[3. 4.]]
Tree Structure:
PyTreeDef({'b1': *, 'extra': (*, *), 'w1': *})
Mapping functions across trees with jax.tree.map
The workhorse of PyTree manipulation is jax.tree.map. It applies a function to every leaf array inside one or more PyTrees, keeping the container architecture untouched.
import jax
import jax.numpy as jnp
params = {
'w': jnp.array([1.0, 2.0, 3.0]),
'b': jnp.array([0.5])
}
# Scale all params by a constant factor
scaled_params = jax.tree.map(lambda x: x * 2.0, params)
print("Original params: ", params)
print("Scaled params: ", scaled_params)
Output
Original params: {'w': Array([1., 2., 3.], dtype=float32), 'b': Array([0.5], dtype=float32)}
Scaled params: {'b': Array([1.], dtype=float32), 'w': Array([2., 4., 6.], dtype=float32)}
jax.tree.map can also accept multiple PyTrees simultaneously, provided they share the same tree-structural definition. This allows for element-wise operations between matching leaves across distinct trees.
Building a Multi-Layer Neural Network Parameter Tree
Now, let's try to build a multi-layer neural network from scratch using pure functions and PyTrees to illustrate how JAX handles multi-layer architectures.
We shall define parameter initialization logic and forward pass for a Multi-Layer Perceptron(MLP)
import jax
import jax.numpy as jnp
def init_mlp_params(layer_sizes, key):
params = []
keys = jax.random.split(key, len(layer_sizes) - 1)
for in_dim, out_dim, k in zip(layer_sizes[:-1], layer_sizes[1:], keys):
w_key, b_key = jax.random.split(k)
scale = jnp.sqrt(2.0 / (in_dim + out_dim))
weight = jax.random.normal(w_key, (in_dim, out_dim)) * scale
bias = jnp.zeros((out_dim,))
params.append({'weights': weight, 'biases': bias})
return params
def mlp_predict(params, x):
activations = x
num_layers = len(params)
for i, layer in enumerate(params):
outputs = jnp.dot(activations, layer['weights']) + layer['biases']
# Apply ReLU for hidden layers and linear output for final layer
activations = jax.nn.relu(outputs) if i < num_layers - 1 else outputs
return activations
root_key = jax.random.key(42)
# network params:
# 3 input features -> 8 hidden units -> 1 output unit
mlp_params = init_mlp_params([3, 8, 1], root_key)
# Test forward pass with dummy input
dummy_x = jnp.array([1.0, -2.0, 0.5])
prediction = mlp_predict(mlp_params, dummy_x)
print("Initialized network param PyTree")
print(f"Number of layers: {len(mlp_params)}")
print(f"Layer 1 weight shape: {mlp_params[0]['weights'].shape}")
print(f"Prediction for single sample: {prediction}")
Output:
Initialized network param PyTree
Number of layers: 2
Layer 1 weight shape: (3, 8)
Prediction for single sample: [0.4792502]
One of JAX's strengths is functional composition, where transformations can be wrapped around each other seamlessly. Suppose we want to compute per gradients accross a batch of data, evaluate loss, and execute the entire operation with compiled XLA performance. We can chain jax.vmap, jax.grad and jax.jit together on functions operating directly on PyTree parameters.
import jax
import jax.numpy as jnp
def compute_loss_single(params, x, y):
""" MSE for a single sample """
prediction = mlp_predict(params, x)
return jnp.squeeze((prediction - y) ** 2)
grad_single_fn = jax.grad(compute_loss_single, argnums=0)
# Vectorize single-sample gradient across batch dimension
per_sample_grads_fn = jax.vmap(grad_single_fn, in_axes=(None, 0, 0))
compiled_per_sample_grads = jax.jit(per_sample_grads_fn)
# Generate batched data
root_key = jax.random.key(101)
x_batch = jax.random.normal(root_key, (5,3))
y_batch = jnp.array([1.0, 0.0, 1.0, 0.5, -1.0])
per_sample_grads = compiled_per_sample_grads(mlp_params, x_batch, y_batch)
print("Per-sample weights gradient shape for Layer 1:", per_sample_grads[0]['weights'].shape)
print("Per-sample biases gradient shape for Layer 1: ", per_sample_grads[0]['biases'].shape)
Output
Per-sample weights gradient shape for Layer 1: (5, 3, 8)
Per-sample biases gradient shape for Layer 1: (5, 8)
By nesting jax.vmap around jax.grad and compiling the pipeline with jax.jit, JAX avoids the standard Python overhead and converts the entire execution into a single GPU/TPU accelerated compute kernel.
Functional Parameter Updates with jax.tree.map
Since JAX arrays are strictly immutable, we cannot perform inplace updates like params -= learning_rate * grad. Instead, model optimization is expressed as a pure function that inputs current parameters and compute gradients returning an updated parameter PyTree.
We achieved this cleanly across arbitrarily nested parameter structures using jax.tree.map
import jax
import jax.numpy as jnp
def compute_batch_loss(params, x_batch, y_batch):
"""MSE averaged over a batch """
predictions = jax.vmap(mlp_predict, in_axes=(None, 0))(params, x_batch)
return jnp.mean((jnp.squeeze(predictions) - y_batch) ** 2)
@jax.jit
def train_step(params, x_batch, y_batch, learning_rate=0.05):
"""Compute batch loss, calculates the batch gradients and applies SGD updates"""
loss, grads = jax.value_and_grad(compute_batch_loss)(params, x_batch, y_batch)
# Perform functional gradient descent update across all PyTree leaves
updated_params = jax.tree.map(
lambda p, g: p - learning_rate * g,
params,
grads
)
return updated_params, loss
# Our training week
root_key = jax.random.key(420)
x_train = jax.random.normal(root_key, (100, 3))
y_train = jnp.sin(x_train[:, 0]) + x_train[:, 1]
current_params = init_mlp_params([3, 16, 1], root_key)
print("Training MLP Model using Functional PyTree Updates...")
for step in range(5)
current_params, loss_val = train_step(current_params, x_train, y_train)
print(f"Step {step + 1} | Batch Loss: {loss_val:.4f}")
Output:
Training MLP Model using Functional PyTree Updates...
Step 1 | Batch Loss: 1.0306
Step 2 | Batch Loss: 0.7979
Step 3 | Batch Loss: 0.6517
Step 4 | Batch Loss: 0.5486
Step 5 | Batch Loss: 0.4708
Notice how jax.tree.map seamlessly traverses all weight and bias arrays across layers in a single line, subtracting the scaled gradient without requiring explicit manual key lookups or conditional checks.
Introduction to Equinox
While managing parameters manually via Python dictionaries and lists is educational, large production neural networks with complex layers can become tedious to organize using raw nested dictionaries.
To handle this, the JAX library offers libraries like Equinox. Equinox introduces class-based module definitions without breaking the purity rules.
In Equinox, model modules are Python dataclasses registered as PyTrees under the hood. Array attributes are treated as PyTree leaves, while non-array attributes like metadata or configuration integers are treated as static structure nodes.
import jax
import jax.numpy as jnp
import equinox as eqx
# Class based module
class EquinoxMLP(eqx.Module):
layers: list
def __init__(self, key):
keys = jax.random.split(key, 2)
self.layers = [
eqx.nn.Linear(in_features=3, out_features=16, key=keys[0]),
eqx.nn.Linear(in_features=16, out_features=1, key=keys[1])
]
def __call__(self, x):
x = jax.nn.relu(self.layers[0](x))
return self.layers[1](x)
# Instantiate Equinox model module
key = jax.random.key(7)
model = EquinoxMLP(key)
# The model instance is a valid PyTree
leaves = jax.tree.leaves(model)
print(f"Total parameter leaf arrays in model: {len(leaves)}")
# Test forward evaluation
sample_input = jnp.array([0.5, 1.2, -0.8])
output = model(sample_input)
print(f"Model output: {output}")
Output
Total parameter leaf arrays in model: 4
Model output: [-0.37182418]
Because an eqx.Module instance is a PyTree, you can pass the model object directly as the parameter argument into jax.grad, jax.jit, and jax.vmap without extra wrapper dictionaries or boilerplate logic.
In Summary
1. A PyTree is any nested structure (dicts, lists, tuples) whose leaves are array data. They allow JAX to process complex parameters in pure functional transformations.
2. You can apply function operations to every leaf in a nested parameter architecture simultaneously while maintaining tree definitions.
3.Advanced combinations like vmap(grad(fn)) work seamlessly on PyTree structures, allowing per-sample gradient computation and fast parallel processing.
4. Model optimization is performed by passing current parameter trees and gradient trees into jax.tree.map, producing updated parameter trees without in-place mutation.
5. Higher-level frameworks like Equinox rely on PyTrees under the hood, enabling clean, class-based neural network code that integrates directly with native JAX transformations.
See you guys in the next one.
