Skip to content

← Flax NNX step 11 of 100

Medium Primitives

NNX Variable vs Param

Jump to the task

Why this matters

NNX variable types label different roles in a model's state. nnx.Param is the usual label for trainable parameters. Other variables can hold running statistics, counters or caches. These labels let you select parts of the state without relying on attribute names.

The type does not update a value by itself. Gradient and optimizer filters determine which variables are differentiated and updated. For example, nnx.Optimizer(..., wrt=nnx.Param) selects parameters. Optax receives the selected values and gradients; it does not inspect an NNX model on its own.

This exercise uses one parameter matrix and two ordinary variables. The forward pass reads only the matrix. The running mean and counter are stored but are not updated or used in the output.

Worked example

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

class MixedState(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.running_mean = nnx.Variable(jnp.zeros((out_features,)))
        self.step_count = nnx.Variable(jnp.array(0.0))

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

model = MixedState(3, 4, rngs=nnx.Rngs(0))
params = nnx.state(model, nnx.Param)
other_variables = nnx.state(model, nnx.All(nnx.Variable, nnx.Not(nnx.Param)))
assert len(jax.tree.leaves(params)) == 1
assert len(jax.tree.leaves(other_variables)) == 2
assert len(jax.tree.leaves(nnx.state(model))) == 3

Each variable in this model holds one array, so each contributes one array leaf. The kernel's twelve numeric entries still form one leaf. Counting array leaves is not the same as counting scalar parameters, and a variable holding a nested value need not contribute just one leaf.

How type filters work

A Param filter includes Param subclasses. A Variable filter includes all Variable subclasses, including Param, BatchStat and Cache. It does not mean “only variables whose exact class is Variable.”

nnx.All(nnx.Variable, nnx.Not(nnx.Param)) selects variables that are not parameters. It can still include BatchStat, Cache or custom non-Param subclasses. This is not an exact-type filter; exact class matching requires an explicit predicate that checks the variable's type.

Filtering selects matching state paths and omits unrelated entries. An attribute named running_mean is not automatically a BatchStat: the wrapper type you construct determines whether that type filter matches it.

Problem

Implement variable_vs_param(seed, x, features):

  1. Define MixedState as above. Initialize the kernel using exactly one rngs.params() key and normal sampling scaled by 1 / sqrt(in_features).
  2. Construct it with nnx.Rngs(int(seed)), input width x.shape[-1] and output width int(features), then compute out by calling it on x.
  3. Count array leaves in nnx.state(model, nnx.Param).
  4. Read the first flattened kernel value with model.kernel[...].reshape(-1)[0].
  5. Return a float32 array containing [leaf_count, kernel_first, out.sum()].

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. Return shape (3,). The first value is 1 for this model.

The numerical fixtures cannot prove that the two unused variables exist or that the leaf count was obtained by filtering. Inspect the state separately when practicing those parts of the exercise.

Common pitfalls

  • Choose wrappers by role. A generic Variable is not selected by a Param filter, even if it happens to contain a floating-point array.
  • A zero-initialized array may be a trainable bias or nontrainable state; its value does not decide its role.
  • Use the array leaf count, not the kernel size or the length of a top-level state mapping. Those measure different things.
  • Keep kernel_first and the output sum as JAX values. Python float conversion breaks tracing; use jnp.stack with a float32 scalar for the count.
  • If applying JIT to this setup function, mark seed and features static because their integer values control initialization and array shapes.