Welcome to the next stage of your JAX journey! Over the past few modules, we've built a solid foundation, starting with JAX's array objects and random number generation, moving through the core function transformations like jit, grad, and vmap, and finally exploring the intricacies of automatic differentiation with techniques like remat. You now have a deep understanding of the fundamental mechanics of JAX.
Now, we'll begin using these tools to build and train neural networks. This module introduces Equinox, a library you specifically wanted to explore. Its design is deeply intertwined with the functional principles of JAX, making it an excellent choice for understanding idiomatic JAX model development.
Today's lesson focuses on the single most important concept in Equinox: its core philosophy. By the end of this hour, you will be able to explain how Equinox represents models as pytrees, a stark and powerful contrast to the approach taken by object-oriented frameworks like PyTorch. This conceptual understanding is the key that unlocks everything else in the library.
1. Models as Data: The Pytree Philosophy
First, let's briefly revisit a concept from Module 3: the pytree. In JAX, a pytree is any nested container of Python objects, typically lists, tuples, and dictionaries. The non-container elements are the "leaves," and the container structure itself is the "treedef."
This concept is not just a convenient data structure; it's central to how JAX operates. Functions in the jax.tree_util library allow you to map operations over the leaves of a pytree while preserving its structure. You've already used this, for instance, to apply optimizer updates.
The core philosophy of Equinox is to take this idea to its logical conclusion: a model is just a pytree.
This is a profound departure from the paradigm you're familiar with in PyTorch.
- In PyTorch, a model (
torch.nn.Module) is a stateful object that has parameters. You interact with these parameters through methods likemodel.parameters()ormodel.state_dict(). The model object itself is largely opaque to PyTorch's core systems; it's a container, but the optimizer and other tools are primarily concerned with the list of tensors it holds. - In Equinox, the model object is the collection of its parameters and logic. The entire
equinox.Moduleinstance—including its weights, biases, submodules, and even static data like activation functions—is structured as a single, unified pytree that is completely transparent to JAX.
This means JAX transformations like jax.grad and jax.jit can be applied directly to the model object itself. There is no need to extract parameters into a separate list. When you take the gradient of a loss function with respect to an Equinox model, you get back a new model instance of the exact same structure, where the leaves that were parameters now contain their corresponding gradients.
2. How It Works: Callable Pytrees
To understand this mechanism, it's insightful to look at the original paper and documentation that introduce Equinox. The key idea is what the author calls a "callable PyTree."
Equinox: neural networks in JAX via callable PyTrees and filtered transformations
This paper by Patrick Kidger introduces the core ideas behind Equinox. It's written for an academic audience but is very clear and concise, which should make it a quick read for you.
First, read the abstract. It perfectly summarizes the problem Equinox solves: bridging the gap between PyTorch's OO-style and JAX's functional paradigm. Next, read Section 1.3 on Contributions. Pay attention to the list of things Equinox does not have—no new abstractions, no complex OO-to-functional transforms. This is the essence of its design philosophy. Finally, and most importantly, read Section 3, Parameterised functions as data. This is the heart of the matter. It explains that Equinox classes are registered as custom PyTree nodes with JAX. This registration tells JAX how to "flatten" a model instance into its leaves (its parameters and static data) and how to "unflatten" them back into a model instance.
The key takeaway is that by implementing tree_flatten and tree_unflatten, an Equinox model makes its internal state—its parameters—completely transparent to JAX. When jax.grad traces a function that uses the model, it can see the model's parameters as leaves of the computational graph, just like any other JAX array.
3. A Practical Example
The "All of Equinox" documentation provides a clear, practical example of this philosophy in action.
This document serves as a high-level user guide. We'll look at the first section, which builds a simple neural network.
Read the section Models as PyTrees. Pay close attention to the following: The NeuralNetwork class inherits from eqx.Module. This base class handles the PyTree registration for you. The attributes self.layers and self.extra_bias are defined. These are the children in the PyTree. The layers attribute is itself a pytree (a tuple of eqx.nn.Linear modules). The loss function takes model as its first argument. jax.grad is applied to loss, differentiating with respect to the model argument. The update step uses jax.tree_util.tree_map, the canonical way to operate on PyTrees. It maps a lambda function over the model and grads trees.
This example beautifully illustrates the entire workflow. The model is created, passed into a loss function as data, differentiated, and updated, all while being treated as a single, coherent pytree.
4. The Challenge: Handling Non-Array Data
A crucial question arises from this design. If the entire model is a pytree, what happens to attributes that aren't JAX arrays? Consider activation functions (like jax.nn.relu), boolean flags, or optimizer configurations. These are valid Python objects but are not differentiable and cannot be handled by a JIT compiler.
Equinox's second key idea is filtering. The library provides tools to dynamically partition a pytree into parts that JAX transformations should see (arrays) and parts they should ignore (static data).
A video from the "Machine Learning & Simulation" channel demonstrates this concept well as part of a full training loop.
Neural Networks in Equinox (JAX DL framework) with Optax
This video provides a practical walkthrough of building and training an MLP in Equinox. It clearly shows the "model as pytree" concept and the role of filtering.
First, watch the model definition from the start. Notice how the SimpleMLP class is defined by its layers, which are eqx.nn.Linear modules. The entire model is just a composition of these array-containing objects. Next, watch the explanation of gradients from this segment. The key insight is here: the presenter explains that eqx.filter_value_and_grad returns a new model instance where the leaves contain the gradients. This is the philosophy in action. Finally, see how filtering is used with the optimizer from this part. The code eqx.filter(model, eqx.is_array) is used to extract only the JAX arrays (the trainable parameters) to initialize the optimizer state. The updates from the optimizer are then applied back to the full model pytree using eqx.apply_updates.
As the video and documentation show, this filtering is typically handled by convenience wrappers like eqx.filter_jit and eqx.filter_grad. These automatically partition the model into dynamic parts (JAX arrays) and static parts (everything else), pass them correctly across the JAX API boundary, and then recombine them. This allows you to write clean, object-oriented code while remaining fully compliant with JAX's functional constraints.
Conclusion
In this lesson, we've focused on the foundational philosophy of Equinox, which is the key to using it effectively. You've seen how it elegantly resolves the tension between the familiar object-oriented style of PyTorch and the strict functional requirements of JAX.
Key Takeaways:
- Models as Pytrees: The core philosophy of Equinox is that a model is a pytree. The model object itself, with all its parameters and submodules, is treated as a single data structure.
- Transparency to JAX: This design makes the model's parameters fully transparent to JAX transformations like
jit,grad, andvmap, which can operate on the model object directly. - No Separate Parameter Dictionaries: Unlike in PyTorch, there's no need for
model.parameters(). The model is its parameters. Gradients are returned as a new model of the same structure. - Filtering for Static vs. Dynamic Data: Equinox uses filtering (e.g.,
eqx.filter_grad) to separate JAX arrays (dynamic parameters) from other Python objects (static configuration) when interfacing with JAX transformations.
With this conceptual framework in place, you are now ready to start building your own models. In our next lesson, we will get hands-on and define our first eqx.Module. We'll explore how to structure it with parameters, submodules, and static fields—the building blocks for the complex architectures you aim to implement.