Transformations on Pytrees

1. Gradients of Pytrees

In PyTorch, automatic differentiation works by tracking operations on Tensor objects that have requires_grad=True and then accumulating gradients in-place via .backward(). JAX takes a different, more functional approach. As the following image illustrates, jax.grad is a transformation that takes a Python function and returns a new function that computes the gradient.

This functional nature is what allows it to integrate seamlessly with pytrees. When you apply jax.grad to a function that accepts a pytree as an argument, JAX understands to compute the gradient with respect to all the differentiable leaves (i.e., floating-point arrays) within that tree. The output is a new pytree, matching the input's structure, containing the gradients.

In pure JAX, you would often use the argnums argument to specify which arguments to differentiate. For example, in a function loss(params, data), jax.grad(loss, argnums=0) would compute gradients with respect to params.

This is demonstrated well in the "Neural Networks in pure JAX" video.

While powerful, this approach requires you to manually separate your parameters from other data. This is where Equinox, which you're focusing on, provides a more streamlined workflow. An Equinox model is a single pytree that mixes trainable parameters (arrays) with static information (like functions or configuration).

How do you take a gradient of such a mixed pytree without errors? Equinox provides filter_grad. It's a wrapper around jax.grad that automatically finds all floating-point arrays in the input pytree, computes gradients for them, and returns a gradient pytree where the gradients for non-differentiable leaves are None. This is exceptionally convenient.

Let's see this in action. The video "Neural Networks in Equinox" provides a great walkthrough.

2. Vectorizing Pytrees with vmap

Just as grad is pytree-aware, so is vmap. When you vmap a function that takes a pytree as an argument, JAX's default behavior is to map over the leading axis of all array leaves while broadcasting all non-array leaves. This simple rule enables powerful patterns.

The most common use is batching. If you have a model and a loss function that works on a single data point, you can make it work on a batch of data by vmapping it.

import jax
import equinox as eqx




# Assume `model` is an eqx.Module and `loss_fn(model, x, y)` works on single examples.
# `x_batch` and `y_batch` have a leading batch dimension.




# vmap the loss function. `in_axes=(None, 0, 0)` means:
# - Don't map over the model (broadcast it)
# - Map over the 0-th axis of x and y
batched_loss_fn = jax.vmap(loss_fn, in_axes=(None, 0, 0))




# Now you can call it on batches
total_loss = batched_loss_fn(model, x_batch, y_batch).mean()

A more advanced and illustrative use of vmap with pytrees is creating and evaluating model ensembles. Imagine you want to train several models with the same architecture but different initializations. You can create them all at once with vmap.

The Equinox documentation has a superb example of this.

This pattern of creating an ensemble by vmapping a constructor is incredibly powerful and not easily replicated in other frameworks. It highlights how treating code and data on an equal footing, combined with pytree-aware transformations, leads to expressive and concise solutions.

3. Composition is Key

The real power of JAX emerges when you start composing these transformations. In the batching example above, we derived the standard training recipe:

1. Start with a loss_fn(model, x, y) that works for a single example.
2. Use vmap to create a batched_loss_fn that works on batches of data.
3. Use grad (or filter_grad) on the batched_loss_fn to get the gradients for the whole batch.

# The canonical training step combines these concepts.

@eqx.filter_jit  # JIT-compile the whole step for speed
def make_step(model, x_batch, y_batch):



    # 1. Define a loss function for a single example
    def loss_fn(m, x, y):
        pred_y = m(x)
        return jax.numpy.mean((pred_y - y) ** 2)




    # 2. Calculate the gradients for the batch
    # grad operates on the output of the mean, which is a scalar.
    # The mean is taken over the output of vmap.
    loss_val, grads = eqx.filter_value_and_grad(loss_fn)(model, x_batch, y_batch)




    # We will cover the update step with Optax in a future lesson
    # new_model = ... update step using grads ...
    
    return loss_val, grads # In a real step, we'd return the new model

While the code looks slightly different (filter_value_and_grad is applied directly to the single-example loss function), JAX is smart enough to push the vmap inside grad where appropriate. What's happening under the hood is effectively grad(mean(vmap(loss))), which is the most efficient way to compute the gradient of the average loss.

Conclusion

In this lesson, you've seen how JAX's core transformations, grad and vmap, are not just for arrays but are deeply integrated with the pytree abstraction. This synergy is what makes the JAX and Equinox programming model so powerful.

Key Takeaways:

• jax.grad and jax.vmap are pytree-aware. They traverse pytree arguments and apply their transformations to the array leaves.
• Equinox's filter_grad and filter_vmap provide convenient wrappers that correctly handle pytrees containing mixed types (arrays, functions, static data) by differentiating/vectorizing arrays and broadcasting everything else.
• Applying grad to a function whose first argument is a model-pytree returns a gradient-pytree of the same structure.
• Applying vmap to a model constructor is an elegant way to create model ensembles, where trainable parameters are stacked along a new axis.
• The standard pattern for batch training involves composing grad and vmap over a loss function defined for a single example.

We've now covered what pytrees are, how to manipulate them, and how JAX's main transformations interact with them. But how does a library like Equinox make its classes behave as pytrees in the first place? In our next lesson, you'll learn how to do this yourself by implementing register_pytree_node_class, giving you the power to make any custom Python object a first-class citizen in the JAX ecosystem.