Skip to content

← Flax NNX step 98 of 100

Medium Primitives

NNX Graphdef vs State Debug

Jump to the task

Why this matters

A model can contain only a few parameter arrays but thousands of individual numbers. Inspecting its state helps you distinguish those counts and find parameters by their paths. This exercise reports three different quantities: the number of array leaves, the sum of their values and the sum of a forward pass's output.

Inspect a two-layer model

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

class TwoLayer(nnx.Module):
    def __init__(self, in_features, hidden, rngs):
        self.l1 = nnx.Linear(in_features, hidden, rngs=rngs)
        self.l2 = nnx.Linear(hidden, hidden, rngs=rngs)

    def __call__(self, x):
        return self.l2(self.l1(x))

model = TwoLayer(3, 4, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)
leaves = jax.tree.leaves(state)
assert len(state) == 2      # top-level entries: l1 and l2
assert len(leaves) == 4     # two kernels and two biases
assert sum(a.size for a in leaves) == 36
total_param_sum = sum(jnp.sum(a) for a in leaves)
assert total_param_sum.shape == ()

The array shapes are (3, 4), (4,), (4, 4) and (4,). Their sizes add up to 36 scalar parameters. Summing their values gives a different statistic, which can be positive, negative or zero.

What these summaries can tell you

The state is nested; counting its top-level entries does not count all its arrays. Tree traversal reaches the arrays inside the Param wrappers. For direct access, state["l1"]["kernel"][...] reads the first kernel.

All state leaves in this model are parameters. A different model may also have counters, running statistics or raw arrays. To inspect only Param variables, use nnx.state(model, nnx.Param). An unfiltered leaf count does not establish what an optimizer will update.

Matching sums do not prove two models are equal. For example, [1, 3] and [2, 2] have the same sum. When verifying a checkpoint, compare structure, shapes, dtypes and corresponding values rather than relying on one total.

The graph definition preserves the model's structure for reconstruction; inspect the state arrays for their shapes and dtypes. Neither the leaf count nor the graph definition alone tells you whether a compiled call can be reused.

Problem

Implement graphdef_vs_state(seed, x, features, hidden):

  1. Build the two-layer model above with input width x.shape[-1] and hidden width int(hidden). Create l1 before l2, passing the same nnx.Rngs(int(seed)) to both constructors. Keep default biases and initializers, and do not insert an activation function.
  2. Compute forward_output_sum = jnp.sum(model(x)).
  3. Split the model and collect the state array leaves with jax.tree.leaves.
  4. Count the leaves and sum all their values to get total_param_sum.
  5. Return [num_leaves, total_param_sum, forward_output_sum] as one array. Keep the sums as JAX scalars and use jnp.stack to combine them with a float32 leaf count. Converting traced sums to Python floats breaks JIT.

There are four leaves here. Calling this model does not modify its state, so inspecting before or after the forward pass gives the same parameter summary. That need not hold for a model with mutable counters or statistics.

Inputs: x is a nonempty, finite float32 vector. seed is a nonnegative integer-valued scalar; hidden is a positive integer-valued scalar. Both are supplied as floats and cast to int. The legacy features argument is accepted but unused; it must not affect the result.

Output: array of shape (3,), containing the leaf count, parameter-value sum and output-value sum. The tests check these numbers, not your choice of inspection API; use the requested split-and-traverse workflow.