Three Levels of Module Nesting

A Flax Linen module can call another module, which can call another. This lets you group a network into reusable parts. In this exercise, three custom modules organize a calculation that contains just two Dense layers and one ReLU.

Follow the calculation

• InnerDense(features) applies one nn.Dense(features=features).
• MiddleBlock(features) calls InnerDense(features) and applies ReLU.
• OuterModel(hidden) calls MiddleBlock(hidden) and applies a final Dense layer with features=x.shape[-1].

After the middle block, x has hidden features. The final Dense therefore maps hidden features to hidden features, regardless of the original input length. There is no final activation, so output values may be negative.

Use these class names and @nn.compact on each __call__ method. Construct and call the child modules inside those methods, as in MiddleBlock(features=self.hidden)(x). Merely constructing a child does not run its computation or initialize the parameters its forward method would create.

Read the parameter paths

For input shape (3,) and hidden=4, jax.tree_util.tree_map(jnp.shape, variables["params"]) gives:

{
    "MiddleBlock_0": {
        "InnerDense_0": {
            "Dense_0": {"kernel": (3, 4), "bias": (4,)}
        }
    },
    "Dense_0": {"kernel": (4, 4), "bias": (4,)}
}

The first Dense belongs to InnerDense_0, inside MiddleBlock_0. The final Dense belongs directly to the outer module. Both are called Dense_0, but their full paths differ, so they have separate parameters. The root module has no additional OuterModel key in this mapping.

Automatic names are local to the containing module, with separate counters for each child class. A module that has no parameters or parameter-bearing children need not have an entry in the params collection. The tree records parameter paths, not every operation in the forward calculation: ReLU adds no parameter entry.

To inspect the first kernel, read:

variables["params"]["MiddleBlock_0"]["InnerDense_0"]["Dense_0"]["kernel"]

Names and paths also affect random initialization. Flattening this network into two Dense calls in one module preserves its mathematical form but can change the initialized weights for the same seed. Keep the nesting and automatic names specified here.

Reuse and selective updates

Calling the same child module instance twice shares its parameters; creating two separate instances gives separate parameter sets. Parameter paths let you identify which arrays to inspect, load or update.

If you later freeze a subtree with Optax, note that optax.masked alone is not a freeze operation: unselected updates pass through unchanged. One option is to append optax.masked(optax.set_to_zero(), freeze_mask) after the optimizer, using boolean True for leaves whose updates should be zero. This exercise does not perform training or require an optimizer.