Skip to content

← Flax NNX step 1 of 100

Medium Primitives

NNX Module Basics

Jump to the task

Why this matters

An NNX module is a Python object that holds parameters and defines how to use them. In this exercise, construction creates the parameters and model(x) computes the output. You can keep that instance and call it again.

Linen uses a different interface: initialize variables once with model.init, then reuse them with model.apply. Linen does not require reinitialization on every forward pass. NNX keeps parameter variables on the model object instead of passing a variables dictionary to each call.

nnx.Param wraps a parameter value and identifies it for parameter filters and the usual gradient/optimizer setup. It does not perform an update by itself. This lesson uses one matrix kernel and no bias, so its output is simply a matrix-vector product written as x @ kernel.

A module constructor does not universally need an RNG argument. Fixed parameters can be initialized without randomness. This exercise uses nnx.Rngs because its kernel is sampled randomly.

Worked example

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

class MyModule(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))
        )

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

model = MyModule(3, 2, rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
first = model(x)
second = model(x)
assert first.shape == (2,)
assert bool(jnp.array_equal(first, second))
assert model.kernel[...].shape == (3, 2)

The RNG container supplies a key when rngs.params() is called in the constructor. Its stream counter advances when a key is requested. The forward method above requests no keys, so calling it again does not reinitialize the kernel.

A new RNG container with the same seed starts the same stream again. Reusing an existing container for another model advances that stream and can produce different initial parameters. Manual JAX key management is also possible; NNX does not forbid it.

Problem

Implement module_basics_forward(seed, x, features):

  1. Define MyModule as above with one nnx.Param kernel, no bias, and a forward method that returns x @ self.kernel.
  2. Draw exactly one key from rngs.params() for normal initialization and scale the kernel by 1 / sqrt(in_features).
  3. Build the module with nnx.Rngs(int(seed)), input width x.shape[-1] and output width int(features). Return model(x).

Return a float32 vector of length features. Numerical fixtures check the output, not whether the kernel was stored as an NNX parameter.

The input x is a nonempty float32 vector. Seed is a nonnegative integer-valued scalar and features is a positive integer-valued scalar; both may arrive as floats. Use the specified key request and normal initializer rather than substituting a default Linear layer initializer.

Common pitfalls

  • Call rngs.params() to obtain a key. The container itself is not a key.
  • Access the parameter array with self.kernel[...] when you need the explicit value. The arithmetic used above also accepts the Param wrapper.
  • Keeping an instance preserves its parameters. Reconstructing it can reset learned values or consume additional RNG keys.
  • Convert the shape argument to an integer. If applying JIT to the setup function, mark seed and features static because they control Python setup.
  • A parameter is model state, not an immutable compile-time constant. Training or an explicit assignment can change it between calls.