Transformers are the dominant architecture in modern deep learning. Let's build a Transformer encoder block in Flax NNX to see how a more complex architecture comes together.
from flax import nnx
import jax
import jax.numpy as jnp
class TransformerBlock(nnx.Module):
"""A single Transformer encoder block."""
def __init__(self, d_model: int, num_heads: int, d_ff: int,
dropout_rate: float, rngs: nnx.Rngs):
# Multi-head self-attention
self.attention = nnx.MultiHeadAttention(
num_heads=num_heads,
in_features=d_model,
qkv_features=d_model,
out_features=d_model,
rngs=rngs,
)
# Feed-forward network
self.ff_linear1 = nnx.Linear(d_model, d_ff, rngs=rngs)
self.ff_linear2 = nnx.Linear(d_ff, d_model, rngs=rngs)
# Layer normalization (pre-norm architecture)
self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
# Dropout
self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)
def __call__(self, x, mask=None):
# Pre-norm self-attention with residual
residual = x
x = self.ln1(x)
x = self.attention(x, mask=mask)
x = self.dropout(x)
x = residual + x
# Pre-norm feed-forward with residual
residual = x
x = self.ln2(x)
x = self.ff_linear1(x)
x = nnx.gelu(x)
x = self.dropout(x)
x = self.ff_linear2(x)
x = self.dropout(x)
x = residual + x
return x
class SmallTransformer(nnx.Module):
"""A small Transformer encoder with positional embeddings."""
def __init__(self, vocab_size: int, d_model: int, num_heads: int,
d_ff: int, num_layers: int, max_len: int,
num_classes: int, rngs: nnx.Rngs):
self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
self.pos_embed = nnx.Embed(max_len, d_model, rngs=rngs)
self.blocks = [
TransformerBlock(d_model, num_heads, d_ff, 0.1, rngs=rngs)
for _ in range(num_layers)
]
self.ln_final = nnx.LayerNorm(d_model, rngs=rngs)
self.classifier = nnx.Linear(d_model, num_classes, rngs=rngs)
def __call__(self, token_ids):
batch_size, seq_len = token_ids.shape
positions = jnp.arange(seq_len)[None, :] # (1, seq_len)
x = self.token_embed(token_ids) + self.pos_embed(positions)
for block in self.blocks:
x = block(x)
x = self.ln_final(x)
x = x[:, 0, :] # CLS token (first position)
return self.classifier(x)
# Create a small Transformer
model = SmallTransformer(
vocab_size=30000, d_model=256, num_heads=8,
d_ff=1024, num_layers=4, max_len=512,
num_classes=5, rngs=nnx.Rngs(42),
)
# Forward pass
token_ids = jnp.ones((2, 128), dtype=jnp.int32)
logits = model(token_ids)
print(f"Output shape: {logits.shape}") # (2, 5)
# Count parameters
num_params = sum(x.size for x in jax.tree.leaves(nnx.state(model)))
print(f"Parameters: {num_params:,}") # ~3.5M💡 Why This Matters
This Transformer block demonstrates that Flax NNX code reads very similarly to PyTorch. The pre-norm architecture (LayerNorm before attention/FFN) is now standard in modern Transformers like GPT and LLaMA. In Track 12, we'll see how to make this memory-efficient with jax.checkpoint/jax.remat and how to use jax.lax.scan to share weights efficiently across layers.
Honesty note: For this simple example, PyTorch's nn.TransformerEncoderLayer would be slightly less code. JAX's advantage shows up when you need to customize training (mixed precision, custom gradient rules, sharding across TPUs) — the functional foundation makes these modifications much easier than monkey-patching a PyTorch module.