Building a CNN: Side-by-Side Comparison

Let's build the same convolutional neural network in both Flax NNX and Equinox, so you can see how the two libraries feel in practice.

Flax NNX Version

from flax import nnx
import jax.numpy as jnp

class CNN_NNX(nnx.Module):
    def __init__(self, num_classes: int, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), padding='SAME', rngs=rngs)
        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), padding='SAME', rngs=rngs)
        self.bn1 = nnx.BatchNorm(32, rngs=rngs)
        self.bn2 = nnx.BatchNorm(64, rngs=rngs)
        self.linear1 = nnx.Linear(64 * 7 * 7, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, num_classes, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

    def __call__(self, x):
        # x: (batch, 28, 28, 1)
        x = nnx.relu(self.bn1(self.conv1(x)))
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nnx.relu(self.bn2(self.conv2(x)))
        x = nnx.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape(x.shape[0], -1)  # flatten
        x = nnx.relu(self.linear1(x))
        x = self.dropout(x)
        x = self.linear2(x)
        return x

# Create and use
model_nnx = CNN_NNX(num_classes=10, rngs=nnx.Rngs(42))
dummy = jnp.ones((4, 28, 28, 1))
out = model_nnx(dummy)
print(f"NNX output shape: {out.shape}")  # (4, 10)

Equinox Version

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

class CNN_Eqx(eqx.Module):
    conv1: eqx.nn.Conv2d
    conv2: eqx.nn.Conv2d
    ln1: eqx.nn.GroupNorm  # GroupNorm instead of BatchNorm
    ln2: eqx.nn.GroupNorm
    linear1: eqx.nn.Linear
    linear2: eqx.nn.Linear
    dropout: eqx.nn.Dropout

    def __init__(self, num_classes: int, *, key):
        k1, k2, k3, k4 = jax.random.split(key, 4)
        self.conv1 = eqx.nn.Conv2d(1, 32, kernel_size=3, padding=1, key=k1)
        self.conv2 = eqx.nn.Conv2d(32, 64, kernel_size=3, padding=1, key=k2)
        self.ln1 = eqx.nn.GroupNorm(groups=8, channels=32)
        self.ln2 = eqx.nn.GroupNorm(groups=8, channels=64)
        self.linear1 = eqx.nn.Linear(64 * 7 * 7, 256, key=k3)
        self.linear2 = eqx.nn.Linear(256, num_classes, key=k4)
        self.dropout = eqx.nn.Dropout(p=0.5)

    def __call__(self, x, *, key=None):
        # x: (28, 28, 1) — Equinox typically works unbatched
        x = jnp.transpose(x, (2, 0, 1))  # (1, 28, 28) — channels first
        x = jax.nn.relu(self.ln1(self.conv1(x)))
        x = eqx.nn.MaxPool2d(kernel_size=2, stride=2)(x)
        x = jax.nn.relu(self.ln2(self.conv2(x)))
        x = eqx.nn.MaxPool2d(kernel_size=2, stride=2)(x)
        x = x.reshape(-1)  # flatten
        x = jax.nn.relu(self.linear1(x))
        x = self.dropout(x, key=key)
        x = self.linear2(x)
        return x

# Create and use — vmap for batching
model_eqx = CNN_Eqx(num_classes=10, key=jax.random.key(42))
dummy_single = jnp.ones((28, 28, 1))
out = model_eqx(dummy_single, key=jax.random.key(1))
print(f"Equinox single output shape: {out.shape}")  # (10,)

# Batch with vmap
keys = jax.random.split(jax.random.key(2), 4)
batch = jnp.ones((4, 28, 28, 1))
out_batch = jax.vmap(model_eqx, in_axes=(0, None))(batch, key=None)
print(f"Equinox batch output shape: {out_batch.shape}")  # (4, 10)