Skip to content

← Flax NNX step 91 of 100

Medium Primitives

NNX Surgery Replace

Jump to the task

Why this matters

Sometimes you want to edit a parameter while keeping the layer and its dimensions. In this exercise, replace every kernel entry in a Linear with the same scalar. Leave the bias unchanged, then apply the layer.

This updates the existing model. If another Python name refers to that model, it sees the change too. Save the output or copy the model before editing if you want to compare before and after.

A complete example

import jax.numpy as jnp
from flax import nnx

model = nnx.Linear(3, 4, rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
before = model(x)
kernel_wrapper = model.kernel

model.kernel[...] = jnp.full_like(model.kernel[...], 0.5)
after = model(x)

assert model.kernel is kernel_wrapper
assert model.kernel[...].shape == (3, 4)
assert jnp.allclose(after, jnp.array([3.0, 3.0, 3.0, 3.0]))

Each output is 1 * 0.5 + 2 * 0.5 + 3 * 0.5 = 3, plus its bias. The default bias is zero. The saved before array still holds the earlier output; updating the model does not retroactively change it.

Keep the parameter wrapper

model.kernel is an nnx.Param. Reading model.kernel[...] gets its array; assigning to model.kernel[...] updates that value while retaining the wrapper. In contrast, assigning a raw array to model.kernel replaces the Param itself. The array may still be part of NNX state, but it no longer matches an nnx.Param filter. That distinction matters for parameter selection.

jnp.full_like uses the template array's shape and dtype by default. The fill value is converted to that dtype, so conversion can round values or overflow. It does not guarantee that arbitrary inputs are representable. Use finite values in the float32 range for this exercise.

This is a parameter-value edit, not a replacement of the Linear submodule. It does not require changing the model's architecture or copying a nested parameter dictionary.

Problem

Implement replace_kernel(seed, x, features, replacement_value):

  1. Build nnx.Linear(x.shape[-1], int(features), rngs=nnx.Rngs(int(seed))) with default initializers and bias.
  2. Fill its kernel with replacement_value, using jnp.full_like and updating the existing Param's array. Leave the bias alone.
  3. Return model(x).

Keep replacement_value as a JAX-compatible scalar. Do not convert it to a Python float: that prevents it from being a dynamic argument under JIT. The seed and feature count configure model construction and must be static if you JIT the whole entrypoint.

In exact arithmetic, each output entry is sum(x) * replacement_value, since every kernel entry is the same and the bias is zero. Floating-point operation order can cause small differences, so compare with a tolerance. This formula explains the output; the exercise still asks you to construct and edit the model. Numerical tests alone cannot verify that workflow.

Inputs: x is a nonempty, finite float32 vector. seed is a nonnegative integer-valued scalar and features is a positive integer-valued scalar, both supplied as floats and cast to int. replacement_value is a finite scalar representable in float32. Inputs are small enough to keep the result finite.

Output: vector of shape (int(features),).