Skip to content

← Flax NNX step 7 of 100

Medium Primitives

NNX Deep Nesting

Jump to the task

Why this matters

Nested modules let you organize a model into smaller pieces. An outer module can delegate to a middle module, which delegates to a layer that owns the actual parameters.

Three levels of nesting do not necessarily mean three learned layers. This example contains one affine calculation followed by ReLU. The outer and middle modules organize that calculation; only InnerLinear owns parameters.

Worked example

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

class InnerLinear(nnx.Module):
    def __init__(self, in_features, out_features, rngs):
        key = rngs.params()
        self.kernel = nnx.Param(
            jax.random.normal(key, (in_features, out_features))
            * (1.0 / jnp.sqrt(in_features))
        )
        self.bias = nnx.Param(jnp.zeros((out_features,)))

    def __call__(self, x):
        return x @ self.kernel + self.bias

class MiddleModule(nnx.Module):
    def __init__(self, in_features, hidden, rngs):
        self.inner = InnerLinear(in_features, hidden, rngs=rngs)

    def __call__(self, x):
        return jax.nn.relu(self.inner(x))

class OuterModule(nnx.Module):
    def __init__(self, in_features, hidden, rngs):
        self.middle = MiddleModule(in_features, hidden, rngs=rngs)

    def __call__(self, x):
        return self.middle(x)

model = OuterModule(4, 8, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)
inner_state = state["middle"]["inner"]
assert inner_state["kernel"][...].shape == (4, 8)
assert inner_state["bias"][...].shape == (8,)
assert len(state) == 1
assert len(jax.tree.leaves(state)) == 2
assert model(jnp.ones(4, dtype=jnp.float32)).shape == (8,)

The attribute path model.middle.inner.kernel corresponds to nested state keys in this model. The kernel and bias are still just two array leaves, even though accessing them passes through three module levels. Models with shared references require additional care; do not assume every arbitrary object graph has a simple one-to-one tree layout.

Randomness through the constructors

Pass the same Rngs instance through all three constructors. Only InnerLinear requests a params key in this example. The outer wrappers do not consume keys just because they exist.

The supplied class expects an Rngs container, not a single random key. Passing a key in its place would not provide a params method. Drawing an unused key in an outer constructor would advance the stream and change the kernel initialized by the inner layer.

Problem

Implement deep_nested_forward(seed, x, hidden):

  1. Define the three classes exactly as organized in the example.
  2. InnerLinear initializes one Param kernel with normal sampling scaled by 1 / sqrt(in_features) and a zero Param bias.
  3. MiddleModule applies ReLU to InnerLinear's output. OuterModule calls MiddleModule and returns that result.
  4. Construct OuterModule with input width x.shape[-1], hidden width int(hidden) and nnx.Rngs(int(seed)), then return its output on x.

Input x is a nonempty finite float32 vector. Hidden is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; both may arrive as floats. Return a float32 vector of length hidden. The complete calculation is relu(x @ kernel + bias).

Common pitfalls

  • Preserve the specified scaled-normal initializer and use exactly one params key. Default nnx.Linear initialization is a different operation.
  • Include the zero bias to preserve the intended parameter structure. Omitting it would leave these initial numerical outputs unchanged, so output tests alone cannot detect that omission.
  • Place ReLU in MiddleModule as specified. A second ReLU would be numerically redundant for these finite values, not an extra learned layer.
  • Do not count module wrappers as parameter arrays. Only the inner kernel and bias contribute leaves here.
  • If JIT-compiling the setup function, mark seed and hidden static because their integer values control initialization and shapes.

Numerical fixtures cannot prove that a solution used three nested modules. Inspect the example's state paths as well as its output when practicing this organization.