Common Layers and Building Blocks

Both Flax NNX and Equinox provide standard neural network layers. Here's a reference for the most common ones:

Flax NNX Layers

from flax import nnx

rngs = nnx.Rngs(0)

# Dense (Linear) layer
dense = nnx.Linear(in_features=256, out_features=128, rngs=rngs)

# Convolution
conv = nnx.Conv(
    in_features=3,        # input channels
    out_features=64,      # output channels
    kernel_size=(3, 3),   # filter size
    strides=(1, 1),
    padding='SAME',
    rngs=rngs,
)

# Embedding
embed = nnx.Embed(
    num_embeddings=10000,  # vocabulary size
    features=256,          # embedding dimension
    rngs=rngs,
)

# Normalization
layer_norm = nnx.LayerNorm(num_features=256, rngs=rngs)
batch_norm = nnx.BatchNorm(num_features=256, rngs=rngs)

# Dropout
dropout = nnx.Dropout(rate=0.1, rngs=rngs)

# Multi-head attention
attention = nnx.MultiHeadAttention(
    num_heads=8,
    in_features=256,
    qkv_features=256,
    out_features=256,
    rngs=rngs,
)

Equinox Layers

import equinox as eqx
import jax

key = jax.random.key(0)
keys = jax.random.split(key, 6)

# Dense (Linear) layer
dense = eqx.nn.Linear(in_features=256, out_features=128, key=keys[0])

# Convolution
conv = eqx.nn.Conv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=3,
    stride=1,
    padding=1,
    key=keys[1],
)

# Embedding
embed = eqx.nn.Embedding(
    num_embeddings=10000,
    embedding_size=256,
    key=keys[2],
)

# Normalization
layer_norm = eqx.nn.LayerNorm(shape=(256,))
# Equinox doesn't have built-in BatchNorm (by design — it requires state)

# Dropout
dropout = eqx.nn.Dropout(p=0.1)

# Multi-head attention
attention = eqx.nn.MultiheadAttention(
    num_heads=8,
    query_size=256,
    key=keys[3],
)

Here's a comparison with PyTorch equivalents for reference:

# PyTorch                        → Flax NNX                → Equinox
# nn.Linear(256, 128)           → nnx.Linear(256, 128)    → eqx.nn.Linear(256, 128)
# nn.Conv2d(3, 64, 3, padding=1)→ nnx.Conv(3, 64, (3,3)) → eqx.nn.Conv2d(3, 64, 3)
# nn.Embedding(10000, 256)      → nnx.Embed(10000, 256)   → eqx.nn.Embedding(10000, 256)
# nn.LayerNorm(256)             → nnx.LayerNorm(256)      → eqx.nn.LayerNorm((256,))
# nn.Dropout(0.1)               → nnx.Dropout(rate=0.1)   → eqx.nn.Dropout(p=0.1)
# nn.MultiheadAttention(256, 8) → nnx.MultiHeadAttention  → eqx.nn.MultiheadAttention