NNX Port: Linen Dense → NNX Linear

Keep the same linear transformation

Linen Dense and NNX Linear both compute x @ kernel + bias in this exercise. Copying the kernel and bias lets an NNX layer reproduce the initialized Linen layer's outputs.

Both use a kernel shaped (in_features, out_features) and a bias shaped (out_features,). No transpose is needed. Linen infers the input width from x during initialization; NNX takes both widths when constructing the layer.

Copy arrays while keeping parameter wrappers

Linen's init returns a variables dictionary. For this layer, the arrays live at variables["params"]["kernel"] and variables["params"]["bias"].

NNX stores those arrays in nnx.Param wrappers. Assign the imported values with linear.kernel[...] = kernel and linear.bias[...] = bias. This keeps the wrappers, so parameter filters and gradient transformations can still identify the weights. Replacing a whole attribute with a bare array changes that structure; whether a forward call errors is not the main issue.

Check shapes and dtypes before transferring weights. Matching names alone do not establish that two arbitrary models are compatible. More complex ports may also need state such as running normalization statistics, configuration settings, and shared-parameter relationships.

A complete example with known weights

Here the input is [1, 2], the kernel is diagonal, and the bias is nonzero. The output is [1*2 + 1, 2*3 - 1], or [3, 5].

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

dense = nn.Dense(features=2)
variables = {"params": {
    "kernel": jnp.array([[2., 0.], [0., 3.]]),
    "bias": jnp.array([1., -1.]),
}}
linear = nnx.Linear(2, 2, rngs=nnx.Rngs(0))
linear.kernel[...] = variables["params"]["kernel"]
linear.bias[...] = variables["params"]["bias"]
x = jnp.array([1., 2.])
assert jnp.allclose(dense.apply(variables, x), jnp.array([3., 5.]))
assert jnp.allclose(linear(x), dense.apply(variables, x))
assert isinstance(linear.kernel, nnx.Param)

For the exercise, initialize Linen with the specified seed rather than using these known weights. The destination NNX layer uses seed+1, but its temporary kernel and bias are both overwritten. Those initial destination values therefore do not determine the result. This is an in-memory transfer; the task does not load a checkpoint or train a model.

Compare outputs, not just sums

The layers should agree numerically after copying their parameters. Use an elementwise tolerance when validating a port, rather than promising bit-for-bit equality across devices, precision settings or compiler versions.

This exercise returns only the two output sums. Equal sums alone do not prove equal outputs: [1, 3] and [2, 2] both sum to four. Inspect the full output arrays while developing your solution, and perform the actual parameter transfer rather than returning one sum twice.