Skip to content

← Flax NNX step 10 of 100

Medium Primitives

NNX Inspecting Model State

Jump to the task

Why this matters

Inspecting model state helps you check which arrays exist and how they are organized. nnx.split(model) returns a graph definition and state. The graph definition records the model structure and static configuration; the state contains values such as kernels, biases and other array state.

A pytree is a nested structure that JAX can flatten into leaves and a structure description. It is not simply a flat list. Here, the leaves are the Linear layer's kernel and bias arrays.

Worked example

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

model = nnx.Linear(4, 6, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)
leaves = jax.tree.leaves(state)

assert state["kernel"][...].shape == (4, 6)
assert state["bias"][...].shape == (6,)
assert len(leaves) == 2
assert sum(array.size for array in leaves) == 30

no_bias = nnx.Linear(4, 6, use_bias=False, rngs=nnx.Rngs(0))
assert len(jax.tree.leaves(nnx.state(no_bias))) == 1

This distinguishes two useful counts. The default layer has two array leaves, containing 24 kernel entries and 6 bias entries. A layer without bias has one array leaf, containing 24 entries.

For this layer, all array state is parameter state. That is not true for every model: running statistics, caches or counters can also be present. Filter with nnx.state(model, nnx.Param) before counting scalar parameter entries in a model that also holds other state.

Reading nested state

A nested model can have state entries containing more State objects. len(state) counts only its immediate keys; jax.tree.leaves(state) visits the leaves recursively. These counts need not agree.

The graph definition and state can be passed to nnx.merge to reconstruct a model. Their existence does not guarantee a single compilation: JIT reuse also depends on array shapes, dtypes and static arguments. Do not cache an old state snapshot without considering later model updates.

Problem

Implement module_state_count(seed, x, hidden):

  1. Construct nnx.Linear(x.shape[-1], int(hidden), rngs=nnx.Rngs(int(seed))) with the default bias enabled.
  2. Split it into graph definition and state.
  3. Count the state's array leaves.
  4. Return a float32 array [leaf_count, hidden, x.shape[-1]].

Input x is a nonempty float32 vector; only its width is used. Hidden is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; both may arrive as floats. Return shape (3,). The first entry is 2 for this default Linear layer. There is no need to run a forward pass or print anything to produce the answer.

Common pitfalls

  • Unpack both values from split. The returned pair is not itself the state.
  • Do not substitute the number of scalar parameters for the array leaf count.
  • Do not assume a bare array is absent from state. Variable wrappers supply roles for filtering; unfiltered state can also contain plain arrays.
  • These counts are Python integers, so converting them to floats is safe. Setting the output array's dtype to float32 is another option.
  • If JIT-compiling this setup function, mark seed and hidden static because their integer values control initialization and shapes.

The output fixtures cannot prove that a solution actually split a model. Use the example's state-shape assertions to practice the intended inspection.