Skip to content

← Flax NNX step 6 of 100

Medium Primitives

NNX Module With Submodule

Jump to the task

Why this matters

A small multilayer perceptron can be built by attaching two Linear layers to an outer NNX module. Each layer owns its kernel and bias; the outer module defines how their outputs are combined.

NNX discovers these submodules through their attributes. Extracting the outer model's state includes the parameters nested under linear1 and linear2. You do not need a separate registration call.

Worked example

import jax
import jax.numpy as jnp
from flax import nnx

class MLP(nnx.Module):
    def __init__(self, in_features, hidden, out, rngs):
        self.linear1 = nnx.Linear(in_features, hidden, rngs=rngs)
        self.linear2 = nnx.Linear(hidden, out, rngs=rngs)

    def __call__(self, x):
        x = self.linear1(x)
        x = jax.nn.relu(x)
        return self.linear2(x)

model = MLP(4, 8, 2, rngs=nnx.Rngs(0))
x = jnp.ones(4, dtype=jnp.float32)
assert model(x).shape == (2,)
assert model.linear1.kernel[...].shape == (4, 8)
assert model.linear2.kernel[...].shape == (8, 2)
params = nnx.state(model, nnx.Param)
assert set(params) == {"linear1", "linear2"}
assert len(jax.tree.leaves(params)) == 4

Construction and randomness

Construct linear1 first, then linear2, passing the same Rngs instance to both. Their initialization consumes keys from the same advancing params stream. Recreating Rngs for each layer or changing the construction order changes the initialization sequence and can change the output.

Keep the default Linear settings to match this exercise: LeCun-normal kernel initialization and a zero bias. A manually sampled normal kernel scaled by input width is not the same initializer. Do not add extra stream calls between constructing the two layers.

The forward method reuses the existing layers. It does not initialize new parameters on each call. Linen also supports reusable parameters, whether submodules are expressed through setup or compact methods; declaring a layer in compact code does not mean it resets on every apply.

Problem

Implement submodule_compose(seed, x, hidden, out):

  1. Define MLP as shown above, creating linear1 before linear2.
  2. Construct it with input width x.shape[-1], hidden width int(hidden), output width int(out) and nnx.Rngs(int(seed)).
  3. Return its output on x: linear1, then ReLU, then linear2.

Input x is a nonempty finite float32 vector. Hidden and out are positive integer-valued scalars; seed is a nonnegative integer-valued scalar. These scalars may arrive as floats. Return a float32 vector of length out.

For a single input, the intermediate shapes are (hidden,) after the first layer and ReLU, then (out,) after the second layer. There is no activation after the final layer in this exercise.

Common pitfalls

  • Build the layers once in the constructor, not in the forward method.
  • ReLU belongs between the layers. Applying it after linear2 would remove negative output values that the fixtures expect.
  • Two attributes pointing to the same layer share its parameters. This exercise creates two distinct Linear instances.
  • Count four parameter array leaves: two kernels and two biases. That is different from counting their scalar entries.
  • If JIT-compiling the setup function, mark seed, hidden and out static because their integer values control initialization and array shapes.

Numerical outputs cannot prove that your code used submodule attributes. Use the example's state assertions to inspect the intended structure.