Skip to content

← Flax NNX step 20 of 100

Medium Primitives

NNX GraphDef Introspection

Jump to the task

Why this matters

To check a layer's dimensions, inspect its arrays. nnx.split(model) gives you a graph definition describing the model's structure and a state mapping containing its dynamic values. You can read array shapes from that state without running a forward pass.

Inspect a Linear layer

import jax
from flax import nnx

model = nnx.Linear(4, 8, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)

kernel_shape = state["kernel"][...].shape
bias_shape = state["bias"][...].shape
assert kernel_shape == (4, 8)
assert bias_shape == (8,)
assert len(jax.tree.leaves(state)) == 2

The kernel has one row per input feature and one column per output feature. The bias has one entry per output feature. This default Linear therefore has two array leaves containing 40 scalar values in total: 32 weights and 8 biases. With use_bias=False, only the kernel contributes an array leaf.

state["kernel"] is a Param wrapper in this example. Indexing it with [...] reads its array. jax.tree.leaves(state) traverses the wrappers and nested mappings to return the array leaves directly.

Structure and array metadata

The graph definition records structure such as module types, attributes, static configuration and shared references. Together with compatible state, it lets nnx.merge(graphdef, state) reconstruct a model.

Do not treat the graph definition as a universal catalogue of array shapes and dtypes. Those properties are available on the state arrays. A Linear also stores configured input and output dimensions, but that does not mean every array's shape is encoded in its graph definition.

A stable graph definition alone does not guarantee JIT cache reuse: array shapes, dtypes and other compilation inputs also matter. You do not need to inspect graph-definition internals or reason about caching for this exercise.

Problem

Implement graphdef_call_static(seed, x, hidden):

  1. Construct nnx.Linear(x.shape[-1], int(hidden), rngs=nnx.Rngs(int(seed))) with its default bias and initializers.
  2. Split the model with nnx.split.
  3. Read the kernel's two dimensions and the bias's single dimension from their arrays in the state.
  4. Return those three dimensions as a float32 array.

The result is [x.shape[-1], int(hidden), int(hidden)]. The exercise is about inspecting the constructed model, even though a shortcut using the inputs alone would produce the same numbers. Numerical tests cannot prove that you split the model or read its state.

Inputs: x is a nonempty float32 vector; its values are unused. seed is a nonnegative integer-valued scalar and hidden is a positive integer-valued scalar, both supplied as floats and cast to int.

Output: float32 array of shape (3,), containing the input width, output width and bias length. No forward pass is required.