Skip to content

← Flax NNX step 4 of 100

Medium Primitives

NNX Multi-Param Module

Jump to the task

Why this matters

A model can contain parameters used in its calculation and other stored state. Here, a kernel and bias produce the output, while a counter is stored but never changed or read by the forward method.

Param and Variable wrappers label these roles for state filters. A Param filter includes the kernel and bias; a Variable filter also includes the counter because Param is itself a Variable subclass.

Worked example

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

class MultiParam(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,)))
        self.count = nnx.Variable(jnp.array(0.0))

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

model = MultiParam(4, 2, rngs=nnx.Rngs(0))
x = jnp.ones(4, dtype=jnp.float32)
assert set(nnx.state(model, nnx.Param)) == {"kernel", "bias"}
assert set(nnx.state(model, nnx.Variable)) == {"kernel", "bias", "count"}
assert model(x).shape == (2,)
assert model.count[...] == 0.0
assert jnp.allclose(model(x), x @ model.kernel[...])

What the wrappers mean

A Param marks a value that is normally selected for training. It does not compute gradients or update itself. Gradient and optimizer filters determine which values those operations use. Likewise, a generic Variable is not an automatic stop-gradient operation.

Bare JAX arrays are not inherently static: NNX can include them in model state. Use the wrappers here to make the intended roles explicit and to support selection by type. Static configuration, such as a layer width, serves a different purpose from changing array state.

Problem

Implement multi_param_forward(seed, x, features):

  1. Define MultiParam as shown above.
  2. Consume exactly one params key for the kernel. Sample a normal array of shape (in_features, out_features) and multiply by 1 / sqrt(in_features).
  3. Initialize bias to a zero vector and count to scalar zero.
  4. Construct the model with input width x.shape[-1], output width int(features) and nnx.Rngs(int(seed)).
  5. Return x @ kernel + bias. Leave count unchanged.

Input x is a nonempty finite float32 vector. Features 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 features.

The zero bias means the initial output equals x @ kernel. This exercise uses the explicit normal initializer above, not the default initializer of nnx.Linear. Matching only the seed is not enough if you change the sampling operation or consume a different key.

Common pitfalls

  • Wrap both kernel and bias in Param so both match the parameter filter.
  • Use a generic Variable for count and do not increment it in this exercise.
  • Do not assume a Variable filter excludes Params; it includes subclasses.
  • Keep the supplied kernel scale. A different initializer changes the expected numerical output even if the architecture is unchanged.
  • If JIT-compiling this setup function, mark seed and features static because their integer values control initialization and shapes.

The output fixtures cannot prove that the unused counter exists or that the arrays have the required wrappers. Run the example's state assertions when practicing those parts of the implementation.