NNX Bridge: State Translation NNX → Linen

Use NNX weights in a Linen layer

An NNX Linear and a Linen Dense can compute the same affine operation: x @ kernel + bias. For these layers, the kernel layouts match: (input_width, output_width), with bias shaped (output_width,). We can therefore read the NNX arrays and arrange them in the variables dictionary expected by Linen. No transpose or new initialization is needed.

This task translates one layer's state in memory. It does not convert an arbitrary model architecture, save a checkpoint or train either layer.

Separate the graph from its state

graphdef, state = nnx.split(linear) separates the module's static structure from its array state. For this Linear, state contains kernel and bias parameter entries. Read their arrays with state["kernel"][...] and state["bias"][...].

graphdef lets NNX reconstruct its module using nnx.merge, but Linen does not consume that graph definition. You provide the compatible Linen layer separately. Its apply method expects a dictionary whose outer keys are collection names:

variables
  params
    kernel: (input_width, output_width)
    bias:   (output_width,)

The params key is required. Passing only the inner kernel/bias dictionary would leave the Dense layer without the parameter collection it expects.

A complete example

Set known weights in NNX, split its state, then apply them through Linen. For x=[1, 2], 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

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

For the exercise, retain the seeded NNX initialization rather than assigning the example's known weights. Do not initialize the Linen layer separately: its arrays are supplied by the translated variables.

What this translation does not preserve automatically

A plain array dictionary does not encode the NNX graph, variable roles or all wrapper metadata. Existing JAX arrays may retain their device placement and sharding, but a dictionary alone is not a portable specification of those choices or of shared parameter relationships.

In a larger model, match parameter paths, shapes and layer configuration. State outside params belongs in its appropriate Linen collection. For example, variables["batch_stats"] is a sibling of variables["params"], not a child inside it. Its leaf names must also match the destination.

To select only NNX parameters, nnx.state(model, nnx.Param) is a filter. When splitting a mixed-state module, account for the other state too, for example nnx.split(model, nnx.Param, ...). A single split filter does not silently discard unmatched state. This task's Linear has no batch statistics, cache or retained RNG state to translate.

Verify the full outputs

Compare the full arrays with a numerical tolerance. Equal sums alone do not prove equivalent outputs: [1, 3] and [2, 2] both sum to four. The two-sum return format is only a compact exercise result. Perform both forwards and the actual translation instead of repeating one sum.