Skip to content

← Flax NNX step 19 of 100

Medium Primitives

NNX Multi State Types

Jump to the task

Why this matters

A model can hold parameters, running statistics, counters and cached values. Each serves a different purpose, so you may want to inspect or update only part of the state. NNX lets you combine filters to select those parts.

This exercise builds a small model with six array leaves and counts three groups. It does not train the model, normalize activations or use the cache. Variable types label the values; the code using those values determines when they change.

Worked example

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

class BatchStat(nnx.Variable):
    pass

class MultiState(nnx.Module):
    def __init__(self, in_features, hidden, rngs):
        key = rngs.params()
        self.kernel = nnx.Param(
            jax.random.normal(key, (in_features, hidden))
            * (1.0 / jnp.sqrt(in_features))
        )
        self.bias = nnx.Param(jnp.zeros((hidden,)))
        self.running_mean = BatchStat(jnp.zeros((hidden,)))
        self.running_var = BatchStat(jnp.ones((hidden,)))
        self.step_count = nnx.Variable(jnp.array(0.0))
        self.cache = nnx.Cache(jnp.zeros((4, hidden)))

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

model = MultiState(3, 5, rngs=nnx.Rngs(0))
remaining = nnx.All(
    nnx.Variable,
    nnx.Not(nnx.Param),
    nnx.Not(BatchStat),
    nnx.Not(nnx.Cache),
)
assert len(jax.tree.leaves(nnx.state(model))) == 6
assert set(nnx.state(model, nnx.Param)) == {"kernel", "bias"}
assert set(nnx.state(model, BatchStat)) == {"running_mean", "running_var"}
assert set(nnx.state(model, remaining)) == {"step_count"}

Each attribute holds one array, so each contributes one leaf. A larger kernel contains more scalar parameters but still contributes one leaf. The cache is also one leaf despite having shape (4, hidden).

What the combined filter means

nnx.All requires every condition to match. Here, a value must be a Variable and must not be a Param, our custom BatchStat, or a Cache. Type filters also match subclasses, so each exclusion removes that type's subclasses too.

For this model, only step_count remains. This is not a general exact-type filter. A new Variable subclass that does not inherit from any excluded type would also match. The name remaining describes that behavior more accurately than “exactly Variable.”

Our custom BatchStat is distinct from the built-in nnx.BatchStat. Both inherit from Variable, but neither inherits from the other. Excluding the custom class would not exclude an added built-in BatchStat.

Problem

Implement multi_state_filter(seed, x, hidden, eps):

  1. Define BatchStat and MultiState as shown above. Initialize the kernel from one params key with normal sampling scaled by 1 / sqrt(in_features).
  2. Construct the model with nnx.Rngs(int(seed)), input width x.shape[-1] and output width int(hidden).
  3. Extract Param state, custom BatchStat state and state matching the remaining filter in three separate calls to nnx.state.
  4. Count each group's array leaves. Return those counts in that order as a float32 array of shape (3,).

The input x is a nonempty float32 vector. Seed is a nonnegative integer-valued scalar and hidden is a positive integer-valued scalar; both may arrive as floats. This exercise uses x only to obtain the input width.

The eps argument is retained in the function signature but is unused. Ignore it: there is no normalization or division by a variance here. Multiplying eps by zero is unnecessary and would produce NaN for an infinite or NaN value.

The result is [2.0, 2.0, 1.0] for every valid model configuration. These output tests cannot establish that your solution created variables or used filters. Use the example's state-path assertions to check those steps.

Common pitfalls

  • A Variable filter includes Param, BatchStat and Cache subclasses. Use the combined filter to exclude the three groups specified here.
  • Initialize running_var to ones as required. Its value is not used in this calculation, so a different initialization would not cause division by zero.
  • A variable label does not automatically configure an optimizer, checkpoint or device layout. Those operations need their own configuration.
  • Count array leaves rather than numeric entries or top-level state keys.
  • Return jnp.array(counts, dtype=jnp.float32). Do not convert eps to a Python float or otherwise include it in the result.
  • When applying JIT to this setup function, make seed and hidden static because their integer values control initialization and array shapes.