Welcome back to our exploration of JAX and Equinox. In the previous lesson, we moved from the abstract concept of an eqx.Module to implementing concrete neural network layers like Linear, LayerNorm, and Dropout. In doing so, we uncovered a crucial aspect of the Equinox philosophy: models are simply pytrees that can contain a mixture of data types. Beyond the trainable JAX arrays for weights and biases, our models now hold static floats (like dropout probability p) and even Python callables (like activation functions).
This flexibility introduces a challenge. Core JAX transformations like jax.grad and jax.jit are designed to operate on JAX arrays. jax.grad needs to know precisely which arrays to differentiate, and jax.jit needs to distinguish between dynamic array arguments and static, non-array arguments to avoid unnecessary recompilations. Today's lesson introduces the solution: a small but powerful set of Equinox utilities—eqx.filter, eqx.partition, and eqx.combine—that allow us to surgically separate the dynamic, trainable arrays from the static configuration of our models.
1. The Core Idea: Filtering Pytrees
Since an eqx.Module is a pytree, the problem of separating parameters from static data becomes a problem of pytree manipulation. We need a way to create a new pytree that has the exact same structure as our model, but where some leaves are kept and others are replaced (typically with None). This is the essence of filtering.
For instance, consider a simple model:
import equinox as eqx
import jax.numpy as jnp
import jax.nn as nn
class Model(eqx.Module):
weight: jnp.ndarray
bias: jnp.ndarray
activation: callable
model = Model(weight=jnp.ones((2, 2)), bias=jnp.zeros(2), activation=nn.relu)
Our goal is to be able to say, "Give me a version of model that only contains the JAX arrays," or "Give me a version that only contains the non-arrays." The functions eqx.filter, eqx.partition, and eqx.combine are our tools for this.
Let's start by looking at the documentation for these functions.
Manipulating PyTrees - Equinox
The official Equinox documentation provides a concise reference for the three key functions we'll be using today. Getting familiar with their signatures is the first step.
Please review the definitions and examples for the three main functions in the "Filtering, partitioning, combining" section: equinox.filter: Understand its basic operation of keeping or replacing leaves based on a filter_spec. Pay attention to the replace argument, which defaults to None. equinox.partition: Note that this is an efficient way to perform eqx.filter and its inverse simultaneously, splitting a pytree into two distinct pieces. equinox.combine: See how it merges two or more pytrees back together. Its logic is simple: for each leaf position, it picks the first non-None value from the input pytrees.
2. The Canonical Pattern for grad and jit
With this toolbox, we can now establish the fundamental pattern for preparing a mixed-type Equinox model for JAX's transformations. The goal is to split the model into two separate pytrees: one containing the arrays we wish to compute gradients for (the "parameters" or "dynamic" part), and one containing everything else (the "static" part).
We then pass these two separate pieces into our JIT-compiled and differentiated function. Inside that function, we immediately recombine them to get our original model back.
The "All of Equinox" guide provides the definitive example of this workflow.
This guide perfectly illustrates why and how to use the partition/combine pattern. It builds a model containing a non-array leaf (an activation function) and then walks through the process of setting up a loss function for differentiation and compilation.
Please read Section 2: Filtering, focusing on the subsection titled "Option 1: use eqx.{partition,combine}". As you read, focus on these four steps: The Model (NeuralNetwork2): The model is defined to include jax.nn.relu—a Python function—directly inside its layers attribute. This is the piece of static data that necessitates filtering. The Partition: The line params, static = eqx.partition(model, eqx.is_array) is the crucial first step. It splits the model into two new NeuralNetwork2 instances. params contains the JAX arrays, with None at the relu leaf. static contains the relu function, with None for all the array leaves. The filter_spec here, eqx.is_array, is a simple callable that returns True if a leaf is a JAX/NumPy array. The Transformed Function: The loss function is defined to accept params and static as its first two arguments. This allows jax.grad to work on params (its first argument by default) and jax.jit to treat static as a non-changing input by specifying static_argnums=1. This prevents recompilation if only the array parameters change. The Combine: The very first line inside loss is model = eqx.combine(params, static). This seamlessly reconstructs the original model, which can then be used as normal. This partition -> transform -> combine sequence is the explicit, manual way of handling mixed-data models in Equinox. It gives you full control and makes the data flow crystal clear. 3. A Practical Application: Initializing an Optimizer This pattern isn't just an academic exercise; it's a practical necessity. A common place you'll see this is when initializing an optimizer from a library like Optax. The optimizer's init function needs to know the shape and dtype of the parameters it will be optimizing to create its own state (e.g., to store momentum values in Adam). It must be provided with only the trainable parameters. This is a perfect scenario for eqx.filter, as we only need to select the parameters, not split them off from the static part.
This partition -> transform -> combine sequence is the explicit, manual way of handling mixed-data models in Equinox. It gives you full control and makes the data flow crystal clear.
3. A Practical Application: Initializing an Optimizer
This pattern isn't just an academic exercise; it's a practical necessity. A common place you'll see this is when initializing an optimizer from a library like Optax. The optimizer's init function needs to know the shape and dtype of the parameters it will be optimizing to create its own state (e.g., to store momentum values in Adam). It must be provided with only the trainable parameters.
This is a perfect scenario for eqx.filter, as we only need to select the parameters, not split them off from the static part.
Neural Networks in Equinox (JAX DL framework) with Optax
Let's see this in action. The following video demonstrates setting up a training loop for an Equinox model using Optax. Pay close attention to how the optimizer state is initialized.
Watch the segment from this part of the video. The presenter instantiates an Optax optimizer and then creates its initial state with opt.init(). The key line is: opt_state = opt.init(equinox.filter(model, equinox.is_array))
Here, eqx.filter is used to create a temporary copy of the model pytree containing only the JAX arrays. This filtered pytree is passed to the optimizer, ensuring the optimizer state is correctly structured to track only the trainable parameters, while gracefully ignoring static components like activation functions or hyperparameters stored in the model. A common refinement is to use eqx.is_inexact_array, as gradients are only defined for floating-point or complex arrays.
Conclusion
Today we've addressed a fundamental challenge of working with Equinox's rich model structures: preparing them for JAX's array-oriented transformations. You've learned how to wield a powerful set of tools to precisely control the flow of data into functions like jax.grad and jax.jit.
Key Takeaways:
- The Problem:
eqx.Modules are pytrees that can contain a mixture of JAX arrays and static Python objects, but JAX transformations primarily operate on arrays. - The Tools:
eqx.filter: Creates a new pytree by keeping only the leaves that satisfy a condition.
eqx.partition: Efficiently splits a pytree into two parts based on a condition—the ones that satisfy it, and the ones that don't.eqx.combine: Merges partitioned pytrees back into a single, complete pytree.
- The Canonical Pattern: The
partition/combinedance is the core workflow for applyingjax.gradandjax.jitto Equinox models. You partition the model outside the transformed function and combine it back inside. - Practical Use: Filtering is essential for tasks like initializing an optimizer, where a downstream function needs to see only the trainable parameters of your model.
This explicit separation of dynamic arrays from static data is a cornerstone of the functional approach championed by JAX and Equinox. It replaces implicit conventions (like PyTorch's .parameters() method) with transparent, composable functions.
In our next lesson, we will see how Equinox builds on this foundation by providing convenience wrappers, eqx.filter_jit and eqx.filter_grad, which automate this common partition/combine pattern and offer a more concise syntax for the most frequent use cases.