Skip to content

← Flax NNX step 3 of 100

Medium Primitives

NNX Call vs Init

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.

The model below only reads its parameters. Repeated calls with the same input therefore give the same result, provided the parameters have not changed. This is repeatability, not mathematical idempotence: we are not claiming model(model(x)) == model(x).

NNX does not make every forward pass pure. A module can update state or consume random keys, as in dropout, so repeated outputs need not always match. Even a model containing only parameters could explicitly modify them in its call method; the behavior depends on the code.

Worked example

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

class MyLinear(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

model = MyLinear(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 call_vs_init_forward(seed, x, features):

  1. Define MyLinear as above. Its kernel uses exactly one rngs.params() key, normal sampling and the scale 1 / sqrt(in_features). Its bias is a zero vector wrapped in nnx.Param.
  2. Build one instance with nnx.Rngs(int(seed)), input width x.shape[-1] and output width int(features).
  3. Call the same instance twice on x and retain both outputs.
  4. Return their sums in order as a float32 array of shape (2,).

The two reported values are equal for this module. The fixtures check these values, not the number of calls or constructions; computing one output and duplicating its sum could pass the numerical checks.

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.