NNX Bridge: Call NNX From Linen

Keep the Linen init/apply interface

bridge.to_linen exposes an NNX-defined module through Linen's usual interface. This is useful when one new layer is written in NNX but the surrounding model still uses Linen variables and init/apply calls.

In this exercise the NNX module holds one biased Linear layer. The wrapper returns the same affine calculation, x @ kernel + bias, while storing the layer's state in Linen collections.

Pass a constructor, then initialize the wrapper

Give to_linen an NNX class or compatible constructor, not an already initialized model. With the defaults used here, the bridge supplies rngs= when constructing the module. Your constructor must accept it. Other model configuration can be passed as constructor arguments through to_linen, or captured by a local class as in the task.

wrapped.init(key, x) returns the full Linen variables dictionary, including its params collection. Pass that entire dictionary to apply, not only variables["params"]. For this module, parameter paths are:

variables["params"]["linear"]["kernel"]  # (input_width, output_width)
variables["params"]["linear"]["bias"]    # (output_width,)

There is no extra nnx key in this example. The linear path comes from the NNX attribute name. Other architectures have different state paths.

A complete example

The example first initializes the wrapper, then supplies known weights to check its forward. The task retains the seeded weights instead.

import jax
import jax.numpy as jnp
from flax import nnx
from flax.nnx import bridge

class MyNNX(nnx.Module):
    def __init__(self, rngs):
        self.linear = nnx.Linear(2, 2, rngs=rngs)

    def __call__(self, x):
        return self.linear(x)

wrapped = bridge.to_linen(MyNNX)
x = jnp.array([1., 2.])
variables = wrapped.init(jax.random.PRNGKey(0), x)
assert variables["params"]["linear"]["kernel"].shape == (2, 2)
known = {"params": {"linear": {
    "kernel": jnp.array([[2., 0.], [0., 3.]]),
    "bias": jnp.array([1., -1.]),
}}}
assert jnp.allclose(wrapped.apply(known, x), jnp.array([3., 5.]))
assert jnp.allclose(jax.jit(wrapped.apply)(known, x), jnp.array([3., 5.]))

Apply uses the supplied state

The bridge constructs an NNX module internally and restores the supplied Linen state before evaluating it. It does not discard your trained parameters in favor of a new random draw on every apply. Keep the constructor's structure compatible with the variables dictionary.

The initialization key is routed through Linen and the bridge. A direct nnx.Linear(..., rngs=nnx.Rngs(seed)) need not produce the same kernel as this wrapped initialization with the same integer seed. Use the prescribed initialization path to reproduce the exercise outputs.

For this deterministic Linear, repeated apply calls with the same variables and x agree and need no apply-time RNG. Stochastic modules and mutable collections require their own RNG and mutable settings; this example does not cover those cases.

To train through Linen, differentiate a loss with respect to its params argument and supply updated variables on later calls. The NNX origin of the layer does not prevent array gradients. This exercise itself performs no optimization or parameter update.