Skip to content

← Flax NNX step 94 of 100

Medium Primitives

NNX Surgery Add Layer

Jump to the task

Why this matters

You can initialize a larger model by copying compatible layers from a smaller one. Existing parameter values fill the shared part of the model; added layers keep their new initialization. This is one way to reuse learned weights, but whether it improves training depends on the task and architecture.

Here both models are freshly initialized. The exercise practices the copying step, not checkpoint loading or training. It builds a second, longer model rather than appending layers to the original object.

A complete example

import jax.numpy as jnp
from flax import nnx

class MLP(nnx.Module):
    def __init__(self, in_features, num_layers, features, rngs):
        layers = []
        width = in_features
        for _ in range(num_layers):
            layers.append(nnx.Linear(width, features, rngs=rngs))
            width = features
        self.layers = nnx.List(layers)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

small = MLP(3, 2, 4, rngs=nnx.Rngs(0))
big = MLP(3, 3, 4, rngs=nnx.Rngs(1000))
for i in range(2):
    big.layers[i].kernel[...] = small.layers[i].kernel[...]
    big.layers[i].bias[...] = small.layers[i].bias[...]

x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
assert jnp.allclose(big(x), big.layers[2](small(x)), rtol=1e-5, atol=1e-6)
assert big.layers[0].kernel is not small.layers[0].kernel

The larger model now applies the smaller model's two layers, then its own final layer. There are no activation functions in this exercise. Despite the class name MLP, the entire chain therefore computes an affine function.

Copy values, preserve separate parameters

Assigning to kernel[...] and bias[...] updates the destination's existing Param wrappers. Do not assign the source layer or its Param wrapper directly: that would share objects rather than copy values into independent parameters. Later updates to a destination Param should not update the source Param.

Copy biases as well as kernels. Both models start with zero biases here, so forgetting the bias copy can still pass numerical tests. In a trained model, biases may have changed and must be transferred too.

Use nnx.List for the layer sequence as shown. It supports indexing and iteration; no blanket claim about Python lists losing state is needed.

What the extra layers change

In general, the larger model will not initially compute the same function as the smaller one. The added random layers transform its output. A square identity kernel with zero bias would preserve the input to an added Linear in this activation-free chain, but that is not the initialization requested here. Parameter copying alone does not promise faster convergence or better performance.

The larger model uses seed + 1000 to initialize its own layers. This makes copying kernels meaningful for the supplied examples. Different seeds do not imply every value differs: default biases are zero for both models, and an input such as the zero vector can hide output differences.

Problem

Implement add_layer_warm_start(seed, x, num_old_layers, num_extra_layers, features):

  1. Define the model above. Its first Linear maps x.shape[-1] to int(features); subsequent layers map int(features) to itself. Construct them in order, sharing one RNG container within each model.
  2. Build small with int(num_old_layers) layers and nnx.Rngs(int(seed)).
  3. Build big with int(num_old_layers) + int(num_extra_layers) layers and nnx.Rngs(int(seed) + 1000). Keep default initializers and biases.
  4. Copy both kernel and bias values from every layer in small to the corresponding layer in big. Keep destination wrappers distinct.
  5. Leave all added layers unchanged and return big(x).

Inputs: x is a nonempty finite float32 vector. seed is a nonnegative integer-valued scalar; num_old_layers and features are positive integer-valued scalars; num_extra_layers is nonnegative. These scalars arrive as floats and must be cast to int. Values are small enough to keep intermediate computations finite. If JIT-compiling the entire entrypoint, treat seed, layer counts and feature count as static construction arguments.

Output: vector of shape (int(features),). With zero extra layers, it matches small(x); otherwise it is the added layers applied to that output. Numerical tests check the final result, not wrapper independence or every copy operation, so follow the full workflow.