NNX Port: Linen Transformer Block → NNX

Transfer a complete encoder block

This exercise copies an initialized Linen encoder block into a matching NNX block. It includes two LayerNorms, four attention projections and two feed-forward layers. The transfer must preserve both the arrays and the computation that uses them.

Use this pre-normalization layout, where the second branch receives the result of the first residual addition:

a = x + attention(ln1(x))
output = a + ff2(relu(ff1(ln2(a))))

Both LayerNorms normalize the last feature axis independently for each token, use epsilon=1e-6, and have trainable scale and bias. Keep their default variance calculation. ff1 maps D features to F, and ff2 maps F back to D. All linear layers include bias. There is no final ReLU or normalization after the second residual addition.

Specify the same attention on both sides

Input x has shape (T, D). With H heads, each head has d=D/H features. Build attention from four separate biased Dense/Linear layers, named q_proj, k_proj, v_proj and out_proj, each mapping D features to D.

1. Project the normalized input to Q, K and V. Reshape each (T, D) result to (T, H, d), then transpose to (H, T, d).
2. For each head, compute Q @ K.T / sqrt(d) and apply softmax over the final key-position axis. Multiply those weights by V.
3. Transpose the (H, T, d) result back to (T, H, d), reshape to (T, D), and apply out_proj.

Attention is unmasked: every token can attend to every input position. There is no dropout, positional encoding, padding mask or cache in this block. These omissions are part of the specified model, not features you should add during the port.

Both Linen and NNX have built-in attention layers, including projections with separate head dimensions. This exercise uses explicit projections to make the parameter mapping easy to inspect. A different implementation may need reshaping and matching configuration; matching names alone is insufficient to establish equivalence.

Map all sixteen parameter arrays

Give Linen submodules explicit names. Under the variables["params"] dictionary, use the following paths, mirrored by NNX attributes:

LayerNorm has scale and bias, not a kernel. Attention parameters live inside the attn subtree. Copy through parameter[...] to preserve nnx.Param wrappers. No transpose is needed for these matching layouts.

A small normalization transfer example

Here both implementations normalize [1, 3] to approximately [-1, 1] before applying scale [2, 3] and bias [1, -1]. The resulting vector is approximately [-1, 2].

import jax.numpy as jnp
import flax.linen as nn
from flax import nnx

variables = {"params": {
    "scale": jnp.array([2., 3.]), "bias": jnp.array([1., -1.])
}}
source = nn.LayerNorm(epsilon=1e-6)
target = nnx.LayerNorm(2, epsilon=1e-6, rngs=nnx.Rngs(0))
target.scale[...] = variables["params"]["scale"]
target.bias[...] = variables["params"]["bias"]
x = jnp.array([[1., 3.]])
assert jnp.allclose(target(x), source.apply(variables, x))
assert jnp.allclose(target(x), jnp.array([[-1., 2.]]), atol=1e-5)

For the complete block, transfer every listed array, even a scale or bias that happens to equal its destination's initial value. Validate the full outputs with a numerical tolerance. Equal sums alone cannot prove equivalence: [1, 3] and [2, 2] both sum to four. Do not replace the requested two forwards with one result reported twice.