Skip to content

← Flax step 93 of 100

Medium Primitives

Param Surgery - Kernel Replace

Jump to the task

Flax Linen keeps a model's parameter arrays in a variables mapping. You can construct a new mapping that replaces one array, then apply the same model with those changed parameters. This exercise replaces a Dense kernel with a constant while preserving its bias.

Locate the kernel

Here the model is a bare nn.Dense(features), not a custom module containing a Dense layer. After variables = model.init(key, x), the paths are:

  • variables["params"]["kernel"], with shape (in_dim, features).
  • variables["params"]["bias"], with shape (features,).

There is no Dense_0 key. That extra level appears when an unnamed Dense is a child of another module.

Replace one array and keep the rest

Use jnp.full_like(old_kernel, replacement_value) to preserve the kernel's shape and dtype. A scalar does not have the matrix shape expected by Dense. Construct a new params collection containing the replacement kernel and the original bias, then put it under the "params" key for model.apply.

new_variables = {
    "params": {
        **variables["params"],
        "kernel": jnp.full_like(variables["params"]["kernel"], replacement_value),
    }
}

This model has only a params collection. A model with other collections would need those preserved too. Treat the original mapping as immutable regardless of whether its container is a regular dictionary or FrozenDict. Rebuilding the changed branch avoids editing the original dictionary in place; unchanged JAX arrays can be shared between the two mappings.

Predict the result

Dense computes x @ kernel + bias. With every kernel entry set to v, each output feature for row n is v * sum(x[n]) + bias[j]. The default bias is zero, so all output columns in a row agree.

For x = [[1, 2, 3], [-1, -2, -3]], features=2 and v=0.5, the output is [[3, 3], [-3, -3]]. Changing the seed does not change this result because the random kernel has been replaced completely.

Problem

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

  1. Construct a bare nn.Dense(int(features)) using default settings.
  2. Initialize it with jax.random.PRNGKey(int(seed)) and x.
  3. Build a new variables mapping with the kernel filled with float(replacement_value), preserving the initialized zero bias.
  4. Apply the model with that mapping and return the output.

x is a finite float32 JAX array of shape (N, in_dim), with both dimensions positive. features is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; either may arrive as a float. replacement_value is finite. Return shape (N, features) and leave x and the original variables unchanged. Perform the parameter replacement and model application rather than only computing the row-sum shortcut.