# Program Transformations and PyTrees in JAX

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.

1.  Nodes(Branches): These are the structural containers e.g. dicts, lists, tuples, namedtuples or custom registered classes
    
2.  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

```python
# 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`.

```python
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:

```plaintext
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`](http://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.

```python
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

```plaintext
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)

```python
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:

```plaintext
Initialized network param PyTree
Number of layers: 2
Layer 1 weight shape: (3, 8)
Prediction for single sample: [0.4792502]
```

### Combining Transformations on PyTrees

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.

```python
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

```plaintext
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`

```python
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:

```plaintext
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`](http://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.

```python
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

```plaintext
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.