NNX Port: Linen MLP → NNX MLP

Match both layers and the activation

This exercise transfers a two-layer MLP from Linen to NNX. Both models must compute the same sequence:

hidden = relu(x @ W1 + b1) output = hidden @ W2 + b2

ReLU appears only between the layers. There is no final activation, so output values may be negative. Copying weights is not enough if the destination applies a different sequence of operations.

Map the parameter names

The two unnamed Dense instances in this Linen compact method receive names Dense_0 and Dense_1. NNX uses the attributes fc1 and fc2. Copy both kernel and bias for each matching pair:

I is the input width, H the hidden width, and O the output width. No transpose is needed. Assign arrays through parameter[...] to keep the nnx.Param wrappers that identify trainable state.

These names belong to this particular Linen definition. Explicit submodule names or a changed construction order can change the paths. Inspect the source parameter tree instead of assuming every MLP uses the same names. More complex ports also need matching configuration, non-parameter state and any parameter-sharing relationships.

A complete example

This example uses known weights to make the calculation visible. The first layer turns [-1, 2] into [0, 3] after ReLU. The second returns [1, 5]. The exercise below uses initialized Linen weights instead.

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

class LinenMLP(nn.Module):
    hidden: int
    out: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.hidden)(x)
        return nn.Dense(self.out)(jax.nn.relu(x))

class NNXMLP(nnx.Module):
    def __init__(self, in_features, hidden, out, rngs):
        self.fc1 = nnx.Linear(in_features, hidden, rngs=rngs)
        self.fc2 = nnx.Linear(hidden, out, rngs=rngs)

    def __call__(self, x):
        return self.fc2(jax.nn.relu(self.fc1(x)))

variables = {"params": {
    "Dense_0": {"kernel": jnp.eye(2), "bias": jnp.array([0., 1.])},
    "Dense_1": {"kernel": 2 * jnp.eye(2), "bias": jnp.array([1., -1.])},
}}
source = LinenMLP(hidden=2, out=2)
target = NNXMLP(2, 2, 2, nnx.Rngs(0))
for name, layer in [("Dense_0", target.fc1), ("Dense_1", target.fc2)]:
    layer.kernel[...] = variables["params"][name]["kernel"]
    layer.bias[...] = variables["params"][name]["bias"]
x = jnp.array([-1., 2.])
assert jnp.allclose(target(x), jnp.array([1., 5.]))
assert jnp.allclose(target(x), source.apply(variables, x))

Linen and NNX do not promise identical initialized weights from the same integer seed: their key handling and naming differ. Copy the arrays from the specified Linen initialization. The NNX model's temporary values are overwritten, so its initial seed does not affect the fully transferred model.

Check the full result

Compare output arrays element by element with a numerical tolerance. Equal sums alone do not prove a correct port: [1, 3] and [2, 2] both sum to four. This exercise retains a two-sum output format, but you should check the individual outputs while developing the transfer. Do not substitute one model's sum twice for the requested two forwards.