Pytrees in Practice: Neural Networks as Pytrees

Let's put everything together and build a simple neural network represented entirely as a pytree of parameters. This is the fundamental pattern that underlies all JAX neural network libraries.

import jax
import jax.numpy as jnp

def init_mlp(key, layer_sizes):
    """Initialize an MLP as a pytree (list of dicts)."""
    params = []
    for i in range(len(layer_sizes) - 1):
        key, k1, k2 = jax.random.split(key, 3)
        fan_in, fan_out = layer_sizes[i], layer_sizes[i + 1]
        # Glorot/Xavier initialization
        std = jnp.sqrt(2.0 / (fan_in + fan_out))
        params.append({
            'weights': jax.random.normal(k1, (fan_in, fan_out)) * std,
            'bias': jnp.zeros(fan_out),
        })
    return params

def mlp_forward(params, x):
    """Forward pass through the MLP pytree."""
    for i, layer in enumerate(params[:-1]):
        x = x @ layer['weights'] + layer['bias']
        x = jax.nn.relu(x)
    # Last layer: no activation
    x = x @ params[-1]['weights'] + params[-1]['bias']
    return x

# Initialize
key = jax.random.key(42)
params = init_mlp(key, [784, 256, 128, 10])

# Check structure
for i, layer in enumerate(params):
    print(f"Layer {i}: weights {layer['weights'].shape}, "
          f"bias {layer['bias'].shape}")
# Layer 0: weights (784, 256), bias (256,)
# Layer 1: weights (256, 128), bias (128,)
# Layer 2: weights (128, 10), bias (10,)

Training This Pytree Model

def loss_fn(params, x, y):
    """Cross-entropy loss for classification."""
    logits = mlp_forward(params, x)
    log_probs = jax.nn.log_softmax(logits)
    return -jnp.mean(jnp.sum(log_probs * y, axis=-1))

@jax.jit
def train_step(params, x, y, lr=0.001):
    """Single training step — everything is pytrees."""
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    # SGD: subtract scaled gradients from every parameter
    new_params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
    return new_params, loss

# Simulate training
x_batch = jax.random.normal(jax.random.key(0), (32, 784))
y_batch = jax.nn.one_hot(jnp.arange(32) % 10, 10)

for step in range(100):
    params, loss = train_step(params, x_batch, y_batch)
    if step % 20 == 0:
        print(f"Step {step}: loss = {loss:.4f}")

Pytree-Aware Operations

Nearly every JAX operation is pytree-aware:

# grad returns a pytree matching the input
grads = jax.grad(loss_fn)(params, x_batch, y_batch)
# grads[0]['weights'].shape == params[0]['weights'].shape ✓

# jit works on functions that take/return pytrees
fast_forward = jax.jit(mlp_forward)
logits = fast_forward(params, x_batch)

# vmap over the batch dimension while keeping params fixed
single_predict = lambda x: mlp_forward(params, x)
batch_logits = jax.vmap(single_predict)(x_batch)

# Compute parameter norms
param_norms = jax.tree.map(jnp.linalg.norm, params)
# param_norms has same structure: list of {'weights': scalar, 'bias': scalar}

# Total L2 regularization
l2_reg = sum(jnp.sum(x ** 2) for x in jax.tree.leaves(params))
print(f"L2 norm: {jnp.sqrt(l2_reg):.2f}")

Compare this with PyTorch's approach:

# PyTorch: parameters live inside the Module object
# class MLP(nn.Module):
#     def __init__(self):
#         self.layer1 = nn.Linear(784, 256)
#     def forward(self, x):
#         return self.layer1(x)
# model = MLP()
# optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# JAX: parameters are a plain data structure, functions are separate
# params = init_mlp(key, [784, 256, 10])
# grads = jax.grad(loss_fn)(params, x, y)
# params = jax.tree.map(lambda p, g: p - lr * g, params, grads)

# Both work — JAX separates data from logic, PyTorch bundles them together