NNX Bridge: Call Linen From NNX

Reuse a Linen layer inside an NNX model

You can keep an existing Linen layer while building its surrounding model in NNX. bridge.ToNNX wraps the Linen module so an NNX parent can store it as a submodule and call it directly.

This exercise wraps a single biased Linen Dense layer. After initialization it computes x @ kernel + bias, returning F features from an input vector of width I. The wrapper exposes its trainable arrays as NNX parameters; you do not manually copy a variables dictionary.

Initialize before the ordinary forward call

Constructing nn.Dense(features=F) does not create its kernel. Linen needs an input shape to infer I and initialize a kernel of shape (I, F) and a bias of shape (F,).

bridge.lazy_init(model, x) runs the model in initialization mode, including the wrapped Dense, and installs the resulting state. In this recipe an ordinary call before lazy_init fails because the Linen parameters are missing. It is not an implicit initialization path.

The input to lazy_init supplies the shape and dtype needed by Dense; this layer's parameter values do not depend on the sample's numerical contents. Other Linen modules may have different initialization behavior. This bridge operation runs the model; do not assume it is a shape-only calculation with no work or state changes.

A complete example

The example initializes the bridge, then assigns known parameters so the forward result is easy to check. The exercise keeps the original seeded parameters instead.

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

class HybridModel(nnx.Module):
    def __init__(self, features, rngs):
        self.linen_dense = bridge.ToNNX(nn.Dense(features), rngs=rngs)

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

model = HybridModel(2, nnx.Rngs(0))
x = jnp.array([1., 2.])
bridge.lazy_init(model, x)
model.linen_dense.kernel[...] = jnp.array([[2., 0.], [0., 3.]])
model.linen_dense.bias[...] = jnp.array([1., -1.])
assert jnp.allclose(model(x), jnp.array([3., 5.]))
assert jnp.allclose(nnx.jit(lambda m, x: m(x))(model, x), model(x))

Once initialized, call the model normally to reuse its parameters. Calling lazy_init again is another initialization pass and can replace those parameters. Do not put it inside each prediction or training step. A different input feature width requires a compatible layer, not merely a new input to the existing initialized Dense.

Parameters and RNG state have different roles

Inspect trainable arrays after initialization with nnx.state(model, nnx.Param). Before initialization, the wrapped Dense has no kernel or bias yet, but the wrapper may already hold RNG state. An empty parameter selection is not the same as an entirely empty model.

The bridge handles RNG keys on behalf of Linen. Use the specified nnx.Rngs(int(seed)) construction and one initialization pass so the exercise reproduces the expected seeded weights. A standalone Linen init using the same integer seed need not produce those same weights.

An ordinary call can advance the wrapper's RNG counters even though this Dense forward uses no random operation and its parameters stay unchanged. Mutable Linen collections and stochastic layers need their own deliberate RNG and mutable-state settings; this Dense example does not establish that every wrapped module is stateless.

For differentiation or JIT, pass the initialized model explicitly to the NNX transformation. The wrapper's parameters remain available to parameter filters and gradients. No optimizer or training step is required here.