Hello, and welcome back.
In our last lesson, we established the fundamental partition/combine pattern. This is the explicit, manual way to prepare Equinox models, which can contain a mix of JAX arrays and other Python objects, for JAX's core transformations like jit and grad. While this pattern offers maximum clarity and control, it can also feel a bit verbose for such a common operation.
Today, we'll build directly on that foundation by exploring the convenience wrappers Equinox provides: eqx.filter_jit and eqx.filter_grad. You will learn how these functions automate the partition/combine workflow, allowing you to write cleaner, more concise training code. By the end of this lesson, you'll be able to assemble a complete, JIT-compiled training step for an Equinox model, the cornerstone of any training loop.
The Abstraction: From Manual to Automatic Filtering
The partition/combine pattern is so common that Equinox provides decorators that handle it for you. These are not introducing new magical capabilities; they are simply well-designed abstractions over the exact logic you learned in the previous lesson.
The "All of Equinox" guide provides a perfect side-by-side comparison of the manual approach and the filtered-transformation approach.
This document directly contrasts the two methods for applying JAX transformations to Equinox models.
Please read Section 2: Filtering. Pay close attention to the code blocks under "Option 1: use eqx.{partition,combine}" and "Option 2: use filtered transformations". Notice that the loss function in Option 2 is much simpler. It takes model as a direct argument, just like you would intuitively write it. The decorators @eqx.filter_jit and @eqx.filter_grad handle the underlying work: They implicitly partition the model into dynamic (arrays) and static (non-arrays) components. They apply the corresponding JAX transformation (jax.jit or jax.grad) only to the dynamic parts. They ensure the static parts are handled correctly (e.g., as static_argnums for jit) to prevent unnecessary recompilation. They manage the recombination of the parts inside the transformed function.
This shift from explicit data manipulation to declarative decorators is a recurring theme. It allows you to focus on the logic of your model and training step, while the library handles the boilerplate required to interface with JAX's functional, array-oriented world.
The core philosophical difference between PyTorch's and JAX's autodiff mechanisms is worth noting here.
eqx.filter_grad is the Equinox-aware version of jax.grad. When you write @eqx.filter_grad def loss(...), you are creating a new function, loss, which, when called, returns the gradients of the original function's array-based leaves. This functional transformation approach is central to JAX's design and is what enables the powerful compositions we will see next.
Assembling a JIT-Compiled Training Step
Now, let's put these decorators into practice to build an efficient, reusable training step. A typical training step involves:
- Calculating the loss for a batch of data.
- Computing the gradients of the loss with respect to the model's trainable parameters.
- Updating the model's parameters using an optimizer.
The goal is to wrap this entire sequence into a single function and JIT-compile it. This is where JAX shines, as the XLA compiler can fuse these operations into a highly optimized kernel that runs on your accelerator with minimal overhead.
eqx.filter_value_and_grad
First, a crucial utility. It's rare that we only want the gradient. We usually also want the loss value itself for logging. JAX provides jax.value_and_grad for this, and Equinox provides its filtered counterpart: eqx.filter_value_and_grad.
Let's watch a short video that introduces this function and sets up the components for our training step.
Neural Networks in Equinox (JAX DL framework) with Optax
This video builds a simple neural network in Equinox and walks through the process of setting up a loss function and optimizer, leading into the creation of a training step.
First, watch how the loss function is defined and then transformed using eqx.filter_value_and_grad in the segment from this point. The presenter explains that this function is conceptually the same as jax.value_and_grad but is tweaked to work correctly with Equinox models containing non-array leaves. Next, watch the setup of the Optax optimizer from here. Notice the line opt.init(equinox.filter(model, equinox.is_array)). This is exactly the use case we discussed in the last lesson: the optimizer needs to see the pytree structure of the parameters it will be optimizing, so we filter the model to pass it only the arrays.
The make_step Function
With the loss/gradient function and the optimizer ready, we can combine them into a single, pure make_step function. This function will encapsulate one full update.
The "nanoJAXGPT" blog post provides an excellent, real-world example of this pattern.
nanoJAXGPT: A pedagogical introduction to JAX/Equinox
This blog post provides a clean, well-commented implementation of a training loop for a GPT-style model. The structure is canonical for JAX/Equinox projects.
First, look at the compute_loss function. Notice it's decorated with @eqx.filter_jit because it's a self-contained computation that can be compiled. Now, read the entire section on the make_step function. This is the heart of today's lesson. Analyze its structure: Signature: It's a pure function. It takes the model and optimizer_state as inputs. Gradients: loss, grads = eqx.filter_value_and_grad(compute_loss)(model, x, y) is the key line. It computes both the loss and the gradients in one go. Optimizer Update: The optimizer.update function takes the gradients and the current optimizer state to compute the parameter updates. Model Update: eqx.apply_updates creates a new model by applying the computed updates to the old one. This respects immutability. Return Values: It returns the new_model and new_optimizer_state, which will be fed into the next call of make_step. Finally, look at the @eqx.filter_jit decorator on make_step. This is the most important part. By compiling this entire function, JAX/XLA can see the whole data flow—from loss calculation to gradient computation to optimizer update—and fuse it into a single, highly efficient program to be executed on the accelerator.
The Performance Payoff
Why is wrapping everything in a JIT-compiled function so important? Because it dramatically reduces the overhead of communication between the Python interpreter (the "host") and the accelerator (the "device"). Without JIT, each operation (matmul, add, etc.) would require a separate dispatch from host to device, which is slow. With JIT, the entire compiled graph of make_step is sent to the device once, and then executed repeatedly with new data, minimizing host-device interaction.
Let's see just how significant this difference is.
Getting Started with JAX: Training a Handwriting Synthesis Transformer
This video provides an excellent walkthrough of building and training a Transformer model, culminating in a powerful demonstration of JIT's impact.
First, watch the detailed breakdown of the train_step function from this section. The logic is identical to the make_step function we just reviewed: it uses filter_value_and_grad, calls the optimizer, applies updates, and is wrapped in @eqx.filter_jit. Now, watch the dramatic performance comparison starting at this timestamp. With @filter_jit enabled, the training loop runs at over 3 iterations/second. With that single line commented out, performance plummets to less than 0.5 iterations/second (2 seconds per iteration). This is not a micro-optimization; it's a fundamental aspect of writing performant JAX code.
Conclusion
In this lesson, you've seen how to graduate from the manual partition/combine pattern to the clean, declarative approach of using eqx.filter_jit and eqx.filter_grad. These wrappers are the idiomatic way to apply JAX's powerful transformations to complex Equinox models. By composing them, you can build a complete, JIT-compiled training step that is both easy to read and highly performant.
Key Takeaways:
eqx.filter_jitandeqx.filter_gradare convenience wrappers that automate thepartition/combinepattern for applyingjax.jitandjax.gradto Equinox models.- The standard way to compute loss and gradients simultaneously is with
eqx.filter_value_and_grad. - A canonical JAX/Equinox training step is a pure function, decorated with
@eqx.filter_jit, that takes the model and optimizer state as input and returns the updated model and optimizer state. - JIT-compiling the entire training step is crucial for performance, as it allows XLA to fuse operations and minimize host-device communication.
We've now covered the core function transformations jit and grad. However, in many of the examples, you saw another transformation being used extensively: jax.vmap. In our next lesson, we will dive deep into vectorization with vmap, exploring how it allows you to write code for a single example and have JAX automatically and efficiently generalize it to entire batches of data.