NNX Coexistence: Parity With Linen

Compare the same affine operation

A useful migration check holds the inputs and parameters fixed, then compares the old and new implementations. This exercise supplies the same kernel and bias arrays to an NNX Linear and a Linen Dense. Both layers compute x @ kernel + bias with matching dimensions.

The result reports each output's sum and the largest absolute difference between corresponding output elements. The difference is a stronger comparison than the sums alone: [1, 3] and [2, 2] have equal sums but differ by one at each position.

Set up one source of parameters

Initialize the NNX Linear with the supplied seed. Read its current kernel and bias through parameter[...], then build the Linen variables:

{"params": {"kernel": kernel_array, "bias": bias_array}}

Do not call Linen init. Separate initialization would test different weights, and equal integer seeds do not guarantee equal draws through the two initialization paths. Both layers use bias, a kernel shaped (input_width, output_width), and a bias shaped (output_width,).

A complete comparison example

This example uses known weights, giving output [3, 5] and sum eight on both sides. The task below retains the seeded NNX weights instead.

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

linear = nnx.Linear(2, 2, rngs=nnx.Rngs(0))
linear.kernel[...] = jnp.array([[2., 0.], [0., 3.]])
linear.bias[...] = jnp.array([1., -1.])
dense = nn.Dense(2)
variables = {"params": {
    "kernel": linear.kernel[...], "bias": linear.bias[...]
}}
x = jnp.array([1., 2.])
a, b = linear(x), dense.apply(variables, x)
difference = jnp.max(jnp.abs(a - b))
report = jnp.stack([a.sum(), b.sum(), difference])
assert jnp.allclose(report, jnp.array([8., 8., 0.]))

# Equal sums can conceal elementwise differences.
other = b + jnp.array([1., -1.])
assert jnp.allclose(a.sum(), other.sum())
assert jnp.allclose(jnp.max(jnp.abs(a - other)), 1.)

What a zero difference establishes

A zero result says the computed output values agree for this input and these settings. It does not prove that the output arrays share storage, that gradients or mutable state match, or that arbitrary Linen and NNX architectures are interchangeable.

Even equivalent implementations can differ slightly across dtypes, devices, compiler versions or precision settings. Use appropriate elementwise tolerances in a broader migration test. Check several inputs, gradients and relevant state separately when those matter. A mismatch can indicate a porting error, but investigate settings and numerical effects before assigning a cause.

Here both layers use the same float32 arrays and default affine operations. The supplied cases expect zero difference. Calculate the difference from both real outputs; do not hard-code zero or replace the second forward with the first result.