In this tutorial, we explore Equinox, a lightweight and elegant neural network library built on JAX, and show how to use it. We begin by understanding how eqx.Module treats models as PyTrees, which makes parameter handling, transformation, and serialization feel simple and explicit. As we move forward, we work through static fields, filtered transformations such as filter_jit and filter_grad, PyTree manipulation utilities, stateful layers such as BatchNorm, and a complete end-to-end training workflow for a toy regression problem. Throughout the tutorial, we focus on writing clear, executable code that demonstrates not only how Equinox works but also why it fits so well into the JAX ecosystem for research and practical experimentation.
!pip install equinox optax jaxtyping matplotlib -q
import jax
import jax.numpy as jnp
import equinox as eqx
import optax
from jaxtyping import Array, Float, Int, PRNGKeyArray
from typing import Optional
import matplotlib.pyplot as plt
import time
print(f"JAX version : {jax.__version__}")
print(f"Equinox version: {eqx.__version__}")
print(f"Devices : {jax.devices()}")
print("\n" + "="*60)
print("SECTION 1: eqx.Module basics")
print("="*60)
class Linear(eqx.Module):
weight: Float[Array, "out in"]
bias: Float[Array, "out"]
def __init__(self, in_size: int, out_size: int, *, key: PRNGKeyArray):
wkey, bkey = jax.random.split(key)
self.weight = jax.random.normal(wkey, (out_size, in_size)) * 0.1
self.bias = jax.random.normal(bkey, (out_size,)) * 0.01
def __call__(self, x: Float[Array, "in"]) -> Float[Array, "out"]:
return self.weight @ x + self.bias
key = jax.random.PRNGKey(0)
lin = Linear(4, 2, key=key)
leaves, treedef = jax.tree_util.tree_flatten(lin)
print("Leaves shapes:", [l.shape for l in leaves])
print("Treedef:", treedef)
print("\n" + "="*60)
print("SECTION 2: Static fields")
print("="*60)
class Conv1dBlock(eqx.Module):
conv: eqx.nn.Conv1d
norm: eqx.nn.LayerNorm
activation: str = eqx.field(static=True)
def __init__(self, channels: int, kernel: int, activation: str, *, key: PRNGKeyArray):
self.conv = eqx.nn.Conv1d(channels, channels, kernel, padding="same", key=key)
self.norm = eqx.nn.LayerNorm((channels,))
self.activation = activation
def __call__(self, x: Float[Array, "C L"]) -> Float[Array, "C L"]:
x = self.conv(x)
x = jax.vmap(self.norm)(x.T).T
if self.activation == "relu":
return jax.nn.relu(x)
elif self.activation == "gelu":
return jax.nn.gelu(x)
return x
key, subkey = jax.random.split(key)
block = Conv1dBlock(8, 3, "gelu", key=subkey)
x_seq = jnp.ones((8, 16))
out = block(x_seq)
print(f"Conv1dBlock output shape: {out.shape}")We set up the full Equinox environment by installing the required libraries and importing JAX, Equinox, Optax, Jaxtyping, Matplotlib, and other essentials. We immediately verify the runtime by printing the JAX and Equinox versions and the available devices, which helps us confirm that our Colab environment is ready for execution. We then begin with the foundations of Equinox by defining a simple Linear module, creating an instance of it, and inspecting its PyTree leaves and structure before introducing a Conv1dBlock that demonstrates how static fields and learnable layers work together in practice.
print("\n" + "="*60)
print("SECTION 3: Filtered transforms")
print("="*60)
class MLP(eqx.Module):
layers: list
dropout: eqx.nn.Dropout
def __init__(self, in_size, hidden, out_size, *, key: PRNGKeyArray):
k1, k2, k3 = jax.random.split(key, 3)
self.layers = [
eqx.nn.Linear(in_size, hidden, key=k1),
eqx.nn.Linear(hidden, hidden, key=k2),
eqx.nn.Linear(hidden, out_size, key=k3),
]
self.dropout = eqx.nn.Dropout(p=0.1)
def __call__(self, x: Float[Array, "in"], *, key: Optional[PRNGKeyArray] = None) -> Float[Array, "out"]:
for layer in self.layers[:-1]:
x = jax.nn.relu(layer(x))
if key is not None:
key, subkey = jax.random.split(key)
x = self.dropout(x, key=subkey)
return self.layers[-1](x)
key, mk = jax.random.split(key)
mlp = MLP(8, 32, 4, key=mk)
@eqx.filter_jit
def forward(model, x, *, key):
return model(x, key=key)
x_in = jnp.ones((8,))
key, fk = jax.random.split(key)
y_out = forward(mlp, x_in, key=fk)
print(f"MLP output: {y_out}")
@eqx.filter_jit
def loss_fn(model: MLP,
x: Float[Array, "B in"],
y: Float[Array, "B out"],
key: PRNGKeyArray) -> Float[Array, ""]:
keys = jax.random.split(key, x.shape[0])
preds = jax.vmap(model)(x)
return jnp.mean((preds - y) ** 2)
grad_fn = eqx.filter_grad(loss_fn)
key, dk = jax.random.split(key)
X = jax.random.normal(dk, (16, 8))
Y = jax.random.normal(dk, (16, 4))
grads = grad_fn(mlp, X, Y, dk)
print(f"Grad of first layer weight: shape={grads.layers[0].weight.shape}, norm={jnp.linalg.norm(grads.layers[0].weight):.4f}")We focus on Equinox’s filtered transformations by building an MLP that includes both linear layers and dropout. We use filter_jit to compile the forward pass while allowing the model to contain non-array fields, and we use filter_grad to compute gradients only for array leaves that should actually participate in learning. By running a forward pass and then evaluating gradients on synthetic data, we see how Equinox cleanly bridges model definition and differentiable computation in a JAX-friendly way.
print("\n" + "="*60)
print("SECTION 4: PyTree manipulation")
print("="*60)
arrays, non_arrays = eqx.partition(mlp, eqx.is_array)
print("Non-array leaves (structure only):", jax.tree_util.tree_leaves(non_arrays))
trainable_filter = jax.tree_util.tree_map(
lambda _: True, mlp
)
trainable_filter = eqx.tree_at(
lambda m: (m.layers[0].weight, m.layers[0].bias),
trainable_filter,
replace=(False, False),
)
trainable, frozen = eqx.partition(mlp, trainable_filter)
print("Frozen params (first layer weight shape):", frozen.layers[0].weight.shape)
print("Trainable first-layer weight is sentinel:", trainable.layers[0].weight)
key, nk = jax.random.split(key)
new_weight = jax.random.normal(nk, mlp.layers[0].weight.shape)
mlp_updated = eqx.tree_at(lambda m: m.layers[0].weight, mlp, new_weight)
print("Updated first-layer weight norm:", jnp.linalg.norm(mlp_updated.layers[0].weight).item())
print("\n" + "="*60)
print("SECTION 5: Stateful layers — BatchNorm with inference mode")
print("="*60)
class BNModel(eqx.Module):
linear1: eqx.nn.Linear
bn: eqx.nn.BatchNorm
linear2: eqx.nn.Linear
def __init__(self, in_f, hidden, out_f, *, key: PRNGKeyArray):
k1, k2 = jax.random.split(key)
self.linear1 = eqx.nn.Linear(in_f, hidden, key=k1)
self.bn = eqx.nn.BatchNorm(hidden, axis_name="batch")
self.linear2 = eqx.nn.Linear(hidden, out_f, key=k2)
def __call__(self, x, state, *, inference: bool = False):
x, state = self.bn(jax.nn.relu(self.linear1(x)), state, inference=inference)
return self.linear2(x), state
key, bk = jax.random.split(key)
bn_model, bn_state = eqx.nn.make_with_state(BNModel)(4, 16, 2, key=bk)
@eqx.filter_jit
def train_step_bn(model, state, x):
def single(x):
return model(x, state)
outs, states = jax.vmap(single, axis_name="batch", out_axes=(0, None))(x)
return outs, states
x_batch = jax.random.normal(key, (8, 4))
preds, bn_state = train_step_bn(bn_model, bn_state, x_batch)
print(f"BNModel output shape: {preds.shape}")We explore PyTree manipulation utilities that make Equinox especially flexible for research workflows. We partition the model into array and non-array parts, create a trainable filter to freeze the first layer, and use tree_at to perform an immutable update on a specific parameter without rewriting the whole model. We then extend the tutorial to stateful layers by defining a BatchNorm-based model, creating both the model and its state, and running a batched training-style pass that returns updated state information.
print("\n" + "="*60)
print("SECTION 6: Full training loop (ResNet MLP on noisy sine)")
print("="*60)
class ResBlock(eqx.Module):
fc1: eqx.nn.Linear
fc2: eqx.nn.Linear
proj: Optional[eqx.nn.Linear]
def __init__(self, size: int, *, key: PRNGKeyArray):
k1, k2 = jax.random.split(key)
self.fc1 = eqx.nn.Linear(size, size, key=k1)
self.fc2 = eqx.nn.Linear(size, size, key=k2)
self.proj = None
def __call__(self, x):
residual = x
x = jax.nn.gelu(self.fc1(x))
x = self.fc2(x)
return jax.nn.gelu(x + residual)
class ResNetMLP(eqx.Module):
embed: eqx.nn.Linear
blocks: list
head: eqx.nn.Linear
def __init__(self, in_size, hidden, out_size, n_blocks, *, key: PRNGKeyArray):
keys = jax.random.split(key, n_blocks + 2)
self.embed = eqx.nn.Linear(in_size, hidden, key=keys[0])
self.blocks = [ResBlock(hidden, key=keys[i+1]) for i in range(n_blocks)]
self.head = eqx.nn.Linear(hidden, out_size, key=keys[-1])
def __call__(self, x):
x = jax.nn.gelu(self.embed(x))
for block in self.blocks:
x = block(x)
return self.head(x)
def make_dataset(n: int, key: PRNGKeyArray):
xk, nk = jax.random.split(key)
x = jax.random.uniform(xk, (n, 1), minval=-1.0, maxval=1.0)
y = jnp.sin(2 * jnp.pi * x) + 0.1 * jax.random.normal(nk, (n, 1))
return x, y
key, dk = jax.random.split(key)
X_train, Y_train = make_dataset(2048, dk)
key, dk = jax.random.split(key)
X_val, Y_val = make_dataset(512, dk)
key, mk = jax.random.split(key)
model = ResNetMLP(1, 64, 1, n_blocks=4, key=mk)
schedule = optax.warmup_cosine_decay_schedule(
init_value=0.0, peak_value=3e-3,
warmup_steps=200, decay_steps=2000
)
optimiser = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(schedule))
opt_state = optimiser.init(eqx.filter(model, eqx.is_array))
@eqx.filter_jit
def train_step(model, opt_state, x, y):
def compute_loss(model, x, y):
preds = jax.vmap(model)(x)
return jnp.mean((preds - y) ** 2)
loss, grads = eqx.filter_value_and_grad(compute_loss)(model, x, y)
updates, opt_state_new = optimiser.update(
grads, opt_state, eqx.filter(model, eqx.is_array)
)
model_new = eqx.apply_updates(model, updates)
return model_new, opt_state_new, loss
@eqx.filter_jit
def evaluate(model, x, y):
preds = jax.vmap(model)(x)
return jnp.mean((preds - y) ** 2)We build a deeper end-to-end learning example by defining a residual block and a ResNetMLP model for a noisy sine regression task. We generate synthetic training and validation datasets, initialize the model, configure a warmup cosine learning rate schedule, and prepare the optimizer state using only the model’s array leaves. We also define the jitted train_step and evaluate functions, which provide the core training and validation mechanics for the full workflow.
BATCH = 128
EPOCHS = 30
steps_per_epoch = len(X_train) // BATCH
train_losses, val_losses = [], []
t0 = time.time()
for epoch in range(EPOCHS):
key, sk = jax.random.split(key)
perm = jax.random.permutation(sk, len(X_train))
X_s, Y_s = X_train[perm], Y_train[perm]
epoch_loss = 0.0
for step in range(steps_per_epoch):
xb = X_s[step*BATCH:(step+1)*BATCH]
yb = Y_s[step*BATCH:(step+1)*BATCH]
model, opt_state, loss = train_step(model, opt_state, xb, yb)
epoch_loss += loss.item()
val_loss = evaluate(model, X_val, Y_val).item()
train_losses.append(epoch_loss / steps_per_epoch)
val_losses.append(val_loss)
if (epoch + 1) % 5 == 0:
print(f"Epoch {epoch+1:3d}/{EPOCHS} "
f"train_loss={train_losses[-1]:.5f} "
f"val_loss={val_losses[-1]:.5f}")
print(f"\nTotal training time: {time.time()-t0:.1f}s")
print("\n" + "="*60)
print("SECTION 7: Save & load model weights")
print("="*60)
eqx.tree_serialise_leaves("model_weights.eqx", model)
key, mk2 = jax.random.split(key)
model_skeleton = ResNetMLP(1, 64, 1, n_blocks=4, key=mk2)
model_loaded = eqx.tree_deserialise_leaves("model_weights.eqx", model_skeleton)
diff = jnp.max(jnp.abs(
jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_array))[0]
- jax.tree_util.tree_leaves(eqx.filter(model_loaded, eqx.is_array))[0]
))
print(f"Max weight difference after reload: {diff:.2e} (should be 0.0)")
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(train_losses, label="Train MSE", color="#4C72B0")
axes[0].plot(val_losses, label="Val MSE", color="#DD8452", linestyle="--")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("MSE")
axes[0].set_title("Training curves")
axes[0].legend()
axes[0].grid(True, alpha=0.3)
x_plot = jnp.linspace(-1, 1, 300).reshape(-1, 1)
y_true = jnp.sin(2 * jnp.pi * x_plot)
y_pred = jax.vmap(model)(x_plot)
axes[1].scatter(X_val[:100], Y_val[:100], s=10, alpha=0.4, color="gray", label="Data")
axes[1].plot(x_plot, y_true, color="#4C72B0", linewidth=2, label="True f(x)")
axes[1].plot(x_plot, y_pred, color="#DD8452", linewidth=2, linestyle="--", label="Predicted")
axes[1].set_xlabel("x")
axes[1].set_ylabel("y")
axes[1].set_title("Sine regression fit")
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("equinox_tutorial.png", dpi=150)
plt.show()
print("\nDone! Plot saved to equinox_tutorial.png")
print("\n" + "="*60)
print("BONUS: eqx.filter_jit + shape inference debug tip")
print("="*60)
jaxpr = jax.make_jaxpr(jax.vmap(model))(x_plot)
n_eqns = len(jaxpr.jaxpr.eqns)
print(f"Compiled ResNetMLP jaxpr has {n_eqns} equations (ops) for batch input {x_plot.shape}")
BATCH = 128
EPOCHS = 30
steps_per_epoch = len(X_train) // BATCH
train_losses, val_losses = [], []
t0 = time.time()
for epoch in range(EPOCHS):
key, sk = jax.random.split(key)
perm = jax.random.permutation(sk, len(X_train))
X_s, Y_s = X_train[perm], Y_train[perm]
epoch_loss = 0.0
for step in range(steps_per_epoch):
xb = X_s[step*BATCH:(step+1)*BATCH]
yb = Y_s[step*BATCH:(step+1)*BATCH]
model, opt_state, loss = train_step(model, opt_state, xb, yb)
epoch_loss += loss.item()
val_loss = evaluate(model, X_val, Y_val).item()
train_losses.append(epoch_loss / steps_per_epoch)
val_losses.append(val_loss)
if (epoch + 1) % 5 == 0:
print(f"Epoch {epoch+1:3d}/{EPOCHS} "
f"train_loss={train_losses[-1]:.5f} "
f"val_loss={val_losses[-1]:.5f}")
print(f"\nTotal training time: {time.time()-t0:.1f}s")
print("\n" + "="*60)
print("SECTION 7: Save & load model weights")
print("="*60)
eqx.tree_serialise_leaves("model_weights.eqx", model)
key, mk2 = jax.random.split(key)
model_skeleton = ResNetMLP(1, 64, 1, n_blocks=4, key=mk2)
model_loaded = eqx.tree_deserialise_leaves("model_weights.eqx", model_skeleton)
diff = jnp.max(jnp.abs(
jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_array))[0]
- jax.tree_util.tree_leaves(eqx.filter(model_loaded, eqx.is_array))[0]
))
print(f"Max weight difference after reload: {diff:.2e} (should be 0.0)")
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(train_losses, label="Train MSE", color="#4C72B0")
axes[0].plot(val_losses, label="Val MSE", color="#DD8452", linestyle="--")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("MSE")
axes[0].set_title("Training curves")
axes[0].legend()
axes[0].grid(True, alpha=0.3)
x_plot = jnp.linspace(-1, 1, 300).reshape(-1, 1)
y_true = jnp.sin(2 * jnp.pi * x_plot)
y_pred = jax.vmap(model)(x_plot)
axes[1].scatter(X_val[:100], Y_val[:100], s=10, alpha=0.4, color="gray", label="Data")
axes[1].plot(x_plot, y_true, color="#4C72B0", linewidth=2, label="True f(x)")
axes[1].plot(x_plot, y_pred, color="#DD8452", linewidth=2, linestyle="--", label="Predicted")
axes[1].set_xlabel("x")
axes[1].set_ylabel("y")
axes[1].set_title("Sine regression fit")
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("equinox_tutorial.png", dpi=150)
plt.show()
print("\nDone! Plot saved to equinox_tutorial.png")
print("\n" + "="*60)
print("BONUS: eqx.filter_jit + shape inference debug tip")
print("="*60)
jaxpr = jax.make_jaxpr(jax.vmap(model))(x_plot)
n_eqns = len(jaxpr.jaxpr.eqns)
print(f"Compiled ResNetMLP jaxpr has {n_eqns} equations (ops) for batch input {x_plot.shape}")We run the complete training loop across multiple epochs, shuffle the data, process mini-batches, and track both training and validation losses over time. We then serialize the trained model with Equinox utilities, reconstruct a matching skeleton model, verify that deserialization restores the weights correctly, and visualize the learned fit and loss curves. Also, we inspect the compiled computation graph using jax.make_jaxpr, which provides a useful debugging and introspection view of how the trained Equinox model is executed under JAX.
In conclusion, we built a strong practical understanding of how Equinox helps us write clean, modular, and JAX-native deep learning code without adding unnecessary abstraction. We saw how to define custom modules, manage static and trainable components, apply filtered transformations safely, work with stateful layers, train a residual MLP, save and reload model weights, and inspect compiled computations. In doing so, we experienced how Equinox gives us the flexibility of raw JAX while still providing the structure needed for modern model development. As a result, we came away with a complete hands-on foundation that prepares us to use Equinox confidently for more advanced machine learning experiments and research workflows.
Check out the Full Codes with Notebook here. Also, feel free to follow us on Twitter and don’t forget to join our 130k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。