Equinox — Models as Pytrees

Equinox, created by Patrick Kidger, takes a different approach: models are just pytrees. An eqx.Module is a frozen dataclass whose fields can be JAX arrays (trainable), other modules (nested pytrees), or static data. There's no special "NNX magic" — everything is a pytree, and standard JAX transformations work directly.

import equinox as eqx
import jax
import jax.numpy as jnp

class MLP(eqx.Module):
    layers: list
    dropout: eqx.nn.Dropout

    def __init__(self, in_dim, hidden_dim, out_dim, *, key):
        k1, k2, k3 = jax.random.split(key, 3)
        self.layers = [
            eqx.nn.Linear(in_dim, hidden_dim, key=k1),
            eqx.nn.Linear(hidden_dim, out_dim, key=k2),
        ]
        self.dropout = eqx.nn.Dropout(p=0.2)

    def __call__(self, x, key=None):
        x = self.layers[0](x)
        x = jax.nn.relu(x)
        x = self.dropout(x, key=key)
        x = self.layers[1](x)
        return x

# Create model — key for parameter initialization
model = MLP(784, 256, 10, key=jax.random.key(0))

# Call it directly
x = jnp.ones((784,))
y = model(x, key=jax.random.key(1))
print(y.shape)  # (10,)

The key Equinox concepts:

Filtered Transformations

Since Equinox models mix arrays (trainable) with non-array data (activation functions, dropout rates), you use eqx.filter_jit and eqx.filter_grad instead of bare jax.jit and jax.grad:

@eqx.filter_jit
@eqx.filter_grad
def compute_loss(model, x, y):
    pred = jax.vmap(model)(x)
    return jnp.mean((pred - y) ** 2)

# filter_grad automatically differentiates only w.r.t. arrays,
# leaving static fields (dropout rate, etc.) untouched
grads = compute_loss(model, x_batch, y_batch)

# Update with optax
import optax
optimizer = optax.adam(1e-3)
opt_state = optimizer.init(eqx.filter(model, eqx.is_array))

# Apply updates
updates, opt_state = optimizer.update(grads, opt_state, model)
model = eqx.apply_updates(model, updates)

Partition and Filter

# Split model into trainable and static parts
params, static = eqx.partition(model, eqx.is_array)
# params: same tree structure, but non-arrays replaced with None
# static: same tree structure, but arrays replaced with None

# Recombine
model = eqx.combine(params, static)

# Freeze specific layers
def freeze_first_layer(model):
    # Use tree_at to target specific parts
    filter_spec = jax.tree.map(lambda _: True, model)
    filter_spec = eqx.tree_at(
        lambda m: m.layers[0],
        filter_spec,
        replace=jax.tree.map(lambda _: False, model.layers[0])
    )
    return filter_spec

When to choose Equinox over Flax NNX: Equinox if you want minimal abstraction (everything is a pytree), you're doing scientific ML, or you prefer PyTorch-style ergonomics. Flax NNX if you're in the Google ecosystem, want the largest community, or need features like nnx.scan for very deep models.