Skip to content

← Flax NNX step 13 of 100

Medium Primitives

NNX State Filter

Jump to the task

Why this matters

A model can store weights, running statistics and counters together. Type filters let you select the variables you need for a particular operation. For example, you might select parameters for an optimizer or running statistics for inspection.

In this exercise, define a custom variable type and use it to select one array from a model's state. The type is a label: it does not compute a running mean, update a counter or configure sharding by itself.

Worked example

This model has four array leaves: a kernel, a bias, a running mean and a counter. Only the kernel and bias participate in its forward calculation.

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

class BatchStat(nnx.Variable):
    pass

class FilterDemo(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.step = nnx.Variable(jnp.array(0.0))

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

model = FilterDemo(4, 6, rngs=nnx.Rngs(0))
total = nnx.state(model)
params = nnx.state(model, nnx.Param)
stats = nnx.state(model, BatchStat)

assert len(jax.tree.leaves(total)) == 4
assert set(params) == {"kernel", "bias"}
assert set(stats) == {"running_mean"}
assert len(jax.tree.leaves(params)) == 2
assert len(jax.tree.leaves(stats)) == 1
assert len(jax.tree.leaves(nnx.state(model, nnx.BatchStat))) == 0

The last line matters: the custom BatchStat and Flax's nnx.BatchStat are different classes despite sharing a name. The custom class inherits from nnx.Variable, not from nnx.BatchStat.

Reading the counts

Each variable here holds one array. The (4, 6) kernel has 24 numeric entries, but it is one array leaf. Changing the layer widths changes the number of scalar parameters, not these leaf counts.

Type filters include subclasses. A Param filter selects parameters and any Param subclasses. A Variable filter selects all four variables in this example, including the parameters and custom statistic. An attribute's name does not determine its type.

Filtering does not update the selected values. An optimizer's configuration determines which variables it updates; checkpoint code determines what it saves. Neither operation is performed in this exercise.

Problem

Implement state_filter_count(seed, x, hidden):

  1. Define the custom BatchStat and FilterDemo classes shown above.
  2. Construct the model using input width x.shape[-1], output width int(hidden) and nnx.Rngs(int(seed)). Use one params key for the kernel and the initialization shown in the example.
  3. Extract all state, Param state and custom BatchStat state in three separate calls to nnx.state.
  4. Count each state's array leaves and return the counts in that order as a float32 array of shape (3,).

Seed is a nonnegative integer-valued scalar and hidden is a positive integer-valued scalar; both may arrive as floats. The input x is a nonempty float32 vector. This function uses its width but does not need to call the model or inspect the values in x.

The result is [4.0, 2.0, 1.0] for every valid input. The output tests therefore cannot establish that your code created the model or used filters. Run the example's state-path assertions while practicing those steps.

Common pitfalls

  • Keep the custom BatchStat separate from nnx.BatchStat. Importing the built-in class under the same name would replace your custom definition.
  • Inheriting from Param would make the custom statistic match Param filters too. Inherit directly from Variable for this exercise.
  • Count array leaves, not numeric entries. Counting top-level state keys only happens to agree here because the model has no nested submodules.
  • Set dtype=jnp.float32 on the returned array. The counts themselves are Python integers; converting those counts to floats is also safe.
  • If wrapping the setup function in JIT, make seed and hidden static because their integer values control initialization and array shapes.