NNX Bridge: Train Mixed Linen+NNX Model

Train two parameter trees together

This exercise takes one SGD step through a small network:

h = relu(x @ W1 + b1) prediction = h @ W2 + b2

The first layer's arrays come from an initialized NNX Linear. The second layer is a Linen Dense evaluated with its variables dictionary. The hidden width is always four. Both layers include bias, and there is no activation after the second layer.

The loss is the mean squared error over the output vector: mean((prediction - y) ** 2). It is not a sum, and there is no extra factor of one half. The task uses one input vector, not a minibatch.

Make the parameter dependence explicit

Read the NNX kernel and bias with parameter[...] and put their arrays in a plain dictionary. Define the forward as a function of this dictionary and the Linen variables dictionary. Compute the first affine operation directly, apply ReLU, then call the Linen layer's apply.

jax.grad(loss_fn, argnums=(0, 1)) returns gradients for both argument trees. Compute both from the same original parameters, then update every leaf with new_parameter = parameter - lr * gradient. These are simultaneous updates; do not recompute the second gradient after updating the first layer.

This array-based exercise does not use a bridge wrapper and does not write the new first-layer arrays back into the original NNX object. It evaluates the updated dictionaries directly. NNX-aware gradients are another approach when you want to train a retained module object.

A complete known-weight step

Here every kernel entry starts at one, both biases start at zero, x=[1] and y=[0]. The four hidden activations are all one, so the prediction is four and the loss is sixteen. Each parameter's gradient is eight. With lr=0.01, the new prediction is 4 * 0.84 * 0.92 - 0.08, or 3.0112.

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

first = {"kernel": jnp.ones((1, 4)), "bias": jnp.zeros((4,))}
second = {"params": {
    "kernel": jnp.ones((4, 1)), "bias": jnp.zeros((1,))
}}
dense = nn.Dense(1)
x, y = jnp.array([1.]), jnp.array([0.])

def loss(first, second):
    h = jax.nn.relu(x @ first["kernel"] + first["bias"])
    return jnp.mean((dense.apply(second, h) - y) ** 2)

g1, g2 = jax.grad(loss, argnums=(0, 1))(first, second)
new_first = jax.tree.map(lambda p, g: p - 0.01 * g, first, g1)
new_second = jax.tree.map(lambda p, g: p - 0.01 * g, second, g2)
assert jnp.allclose(loss(first, second), 16.)
assert jnp.allclose(loss(new_first, new_second), 3.0112 ** 2)

The exercise below uses seeded initialization instead of these known arrays. Use an NNX Linear with seed, then initialize the Linen Dense with seed+1 on a zero vector of length four. Keep that order and those initialization paths to reproduce the expected values.

A lower loss is not guaranteed

A sufficiently large step can increase the loss. A zero learning rate leaves it unchanged, and zero gradients can leave some parameters unchanged even when lr is positive. Compute the specified SGD step and report its actual loss; do not retry with a smaller rate or force a decrease. A decrease alone does not prove that both trees were updated.

Both gradients must include the chain rule through the second layer and ReLU. Where the first affine output is negative, ReLU blocks its gradient. JAX uses a zero derivative at exactly zero for jax.nn.relu.