Building Custom Equinox Modules

1. The Anatomy of an eqx.Module

At its heart, an eqx.Module is a Python class that inherits from equinox.Module. This base class automatically handles the boilerplate for registering your custom class as a JAX pytree node. This means JAX will know how to flatten your model into its constituent arrays and static data, and how to unflatten them back into a model instance.

The structure is reminiscent of Python's dataclasses. You define the fields of your module as class attributes with type hints. These fields become the children of the module's pytree.

The three key components you'll define are:

1. Fields (Class Attributes): Declared with type hints at the class level. These define the structure of your model's pytree.
2. __init__ method: The constructor where you create and assign the initial values for the fields. This is where you'll initialize weights, biases, and any submodules.
3. __call__ method: The forward pass. This defines what the model does when it's called like a function.

Let's start by looking at a concrete example from the Equinox documentation, which we touched upon last time.

This NeuralNetwork class perfectly illustrates the concepts of parameters and submodules.

• Parameters: self.extra_bias is a parameter. It's a raw jax.Array that is part of the model's trainable state.
• Submodules: self.layers is a tuple containing eqx.nn.Linear objects. Each eqx.nn.Linear is itself an eqx.Module with its own parameters (weight and bias). This demonstrates the compositional nature of Equinox: models are nested pytrees of other models and parameters.

When JAX flattens an instance of this NeuralNetwork, it recursively traverses this structure, collecting all the JAX arrays from extra_bias and from the weights and biases inside each of the Linear layers.

2. Defining a Module from Scratch

To make this more concrete, let's watch a practical demonstration. The following video walks through the process of creating a simple MLP, defining its layers in __init__ and its forward pass in __call__. This is a very common pattern and seeing it coded live is instructive.

This process directly mirrors the structure of PyTorch nn.Modules, yet the underlying mechanism is entirely different. In PyTorch, nn.Linear registers its parameters with the parent module behind the scenes. In Equinox, the eqx.nn.Linear object is simply a value assigned to an attribute, and the pytree mechanism makes it visible to JAX.

3. Static vs. Dynamic Fields

So far, all the fields in our examples (layers, extra_bias) have contained JAX arrays. But what if we want to include attributes that are not arrays, such as an activation function, a string, or a boolean flag? These are what we call static fields.

Equinox handles this transparently. Any field that is not a JAX array (or a container of JAX arrays) is considered static. When you use a filtered transformation like eqx.filter_grad or eqx.filter_jit, Equinox automatically partitions the model pytree.

• The dynamic part (all JAX arrays) is passed to the underlying JAX transformation (jax.grad or jax.jit). These are the parameters that can be differentiated or are involved in compiled computation.
• The static part (everything else) is treated as compile-time constants. The JIT compiler can specialize the compiled code based on the values of these static fields.

Let's modify our NeuralNetwork to include a static field. Instead of hard-coding jax.nn.relu in the __call__ method, we can store it as an attribute.

import jax
import equinox as eqx
from typing import Callable

class FlexibleNetwork(eqx.Module):
    layers: list
    activation: Callable

    def __init__(self, key):
        key1, key2 = jax.random.split(key)
        self.layers = [
            eqx.nn.Linear(2, 8, key=key1),
            eqx.nn.Linear(8, 2, key=key2)
        ]
        self.activation = jax.nn.relu

    def __call__(self, x):
        x = self.layers[0](x)
        x = self.activation(x)
        x = self.layers[1](x)
        return x

In this FlexibleNetwork:

• self.layers is a list of submodules containing dynamic JAX arrays (weights and biases).
• self.activation is a static field. Its value is a Python function (jax.nn.relu), which is not a JAX array.

When we use eqx.filter_grad on a loss function involving this model, Equinox will identify the weights and biases in self.layers as the dynamic parts to differentiate, and it will treat self.activation as a static part of the model's structure.

The "All of Equinox" document illustrates this by placing jax.nn.relu directly into the list of layers.

Whether you store a static field as a separate attribute (like self.activation) or embed it within a collection (like self.layers) is a design choice. The former is often clearer for distinct configuration options, while the latter can be convenient for building sequential models.

Conclusion

In this lesson, you've learned the practical mechanics of building your own neural network modules in Equinox. We moved from the abstract idea of "models as pytrees" to the concrete implementation of eqx.Module subclasses.

Key Takeaways:

• eqx.Module Structure: A custom module inherits from eqx.Module and is typically defined by its class-level field declarations, an __init__ method for initialization, and a __call__ method for the forward pass.
• Parameters and Submodules: These are fields that contain JAX arrays, either directly (jax.Array) or nested within other eqx.Module instances. They constitute the dynamic, trainable part of your model.
• Static Fields: These are fields containing any non-array Python object, such as functions, booleans, or strings. They are treated as fixed configuration by JAX's transformations.
• Composition: eqx.Modules are highly compositional. Complex models are built by nesting simpler modules, forming a clean, hierarchical pytree structure.

You now have the tools to define the building blocks of any neural network. In the next lesson, we will focus on how to work with the model you've created. Specifically, you'll learn how to use Equinox's filter utilities like eqx.partition and eqx.combine to precisely separate the trainable parameters from the static structure for tasks like optimization and state management.