NNX Bridge: Shared Params Across nnx and Linen

Use one kernel in two forward passes

NNX Linear and Linen Dense both use a kernel shaped (input_width, output_width). With bias disabled, each computes x @ kernel. This exercise supplies the NNX layer's kernel array to Linen, evaluates both layers on the same input, and adds their outputs.

The expected result is 2 * (x @ kernel). That equation describes the numerical output; it does not prove that an implementation evaluated both layers or shared an array. Follow the requested two-forward recipe.

Share the array for the current call

linear.kernel is an nnx.Param wrapper. Read its current JAX array with linear.kernel[...], then place that array at variables["params"]["kernel"] for Linen. You do not need to initialize the Dense layer separately, because you already provide its parameter.

The wrappers differ, while the affine operation and kernel layout agree. This compatibility is specific to the chosen layers and settings; arbitrary model ports may require different paths, layouts or state. Both layers must use use_bias=False for this task.

A complete example

With x=[1, 2] and a diagonal kernel [2, 3], each layer returns [2, 6]. Their sum is [4, 12].

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

linear = nnx.Linear(2, 2, rngs=nnx.Rngs(0), use_bias=False)
linear.kernel[...] = jnp.array([[2., 0.], [0., 3.]])
dense = nn.Dense(features=2, use_bias=False)
variables = {"params": {"kernel": linear.kernel[...]}}
x = jnp.array([1., 2.])
combined = linear(x) + dense.apply(variables, x)
assert jnp.allclose(combined, jnp.array([4., 12.]))

# Assigning a new array to NNX does not update the existing dictionary.
linear.kernel[...] = 2 * linear.kernel[...]
assert jnp.allclose(linear(x), jnp.array([4., 12.]))
assert jnp.allclose(dense.apply(variables, x), jnp.array([2., 6.]))

# Rebuild the dictionary from the current parameter to use the update.
variables = {"params": {"kernel": linear.kernel[...]}}
assert jnp.allclose(dense.apply(variables, x), linear(x))

The exercise retains the seeded kernel instead of replacing it with this example's known values.

Equal values are not a persistent parameter tie

JAX arrays are immutable. Replacing an NNX parameter's array does not rewrite a previously constructed Linen dictionary. Sharing the current array therefore does not automatically synchronize future optimizer updates, cached compiled functions or separately maintained models.

For training with one canonical NNX parameter, build the Linen variables from that parameter inside the differentiated forward. Pass the NNX module explicitly to an NNX gradient transformation. Both uses can then contribute to the gradient of that one parameter. If you instead pass two separate parameter arguments with equal values, you have not thereby tied their gradients or updates.

This exercise performs no optimization and claims no memory-allocation guarantee. Its purpose is to show how compatible layers can read the same current parameter array through different APIs.