Skip to content

← Flax NNX step 14 of 100

Medium Primitives

NNX State Transform

Jump to the task

Why this matters

Sometimes you want to change selected parameter arrays while preserving the rest of a model. For example, zeroing one layer's kernel lets you inspect what remains when that layer no longer uses its input.

NNX state is a pytree with variable wrappers and nested structure. JAX's tree utilities can visit its array leaves. A path-aware mapping function also receives each leaf's location, so it can decide which arrays to edit.

This exercise zeros a Linear layer's kernel while keeping its bias at one. For finite inputs, its output is then an array of ones: x @ 0 + 1.

Worked example

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

model = nnx.Linear(4, 6, rngs=nnx.Rngs(0))
model.bias[...] = jnp.ones((6,), dtype=jnp.float32)
graphdef, state = nnx.split(model)

def zero_if_kernel(path, leaf):
    has_kernel_key = any(
        isinstance(part, jax.tree_util.DictKey) and part.key == "kernel"
        for part in path
    )
    return jnp.zeros_like(leaf) if has_kernel_key else leaf

new_state = jax.tree_util.tree_map_with_path(zero_if_kernel, state)
merged = nnx.merge(graphdef, new_state)
x = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)

assert jnp.array_equal(merged.kernel[...], jnp.zeros((4, 6)))
assert jnp.array_equal(merged.bias[...], jnp.ones(6))
assert jnp.array_equal(merged(x), jnp.ones(6))

The callback returns a replacement array for the kernel and returns other leaves unchanged. zeros_like preserves the kernel's shape and dtype. The mapped state retains the structure needed by merge.

Match path components, not substrings

A JAX path contains structured entries such as DictKey, SequenceKey and GetAttrKey. This example checks for a dictionary key equal to kernel. It does not assume a particular name for the Variable's internal value field.

Converting the whole path to a string and searching for the substring kernel is too broad: it could also match kernel_scale or a parent named kernel_block. Exact component matching avoids those accidental matches. A parent key that is exactly kernel would still match this predicate, so inspect the full path before reusing it on a different model structure.

Plain tree_map is useful when the same rule applies to every leaf. tree_map_with_path is useful here because the selection depends on the leaf's location. Type filtering alone would not separate this layer's kernel and bias: both are Params.

Problem

Implement state_zero_kernel(seed, x, features):

  1. Construct nnx.Linear(x.shape[-1], int(features), rngs=nnx.Rngs(int(seed))) with its default initializer and bias configuration.
  2. Set the bias array to float32 ones using model.bias[...].
  3. Split the model into its graph definition and state.
  4. Map over state with the path-aware callback shown above. Zero the kernel and preserve the bias.
  5. Merge the transformed state with the graph definition and return its output on x.

The input x is a nonempty finite float32 vector. Seed is a nonnegative integer-valued scalar and features is a positive integer-valued scalar; both may arrive as floats. Return a float32 vector of length features.

Every fixture expects ones. These output tests cannot prove that a solution used state mapping or preserved the original model. Use the example's parameter assertions to check the intended intermediate result.

Common pitfalls

  • Zeroing every leaf also zeros the bias. Select the kernel explicitly.
  • Use zeros_like, not multiplication by zero: multiplying NaN or infinity by zero does not produce a finite zero.
  • A tree mapping returns a pytree, not a callable model. Merge it to create a model, or use nnx.update when intentionally updating an existing model.
  • This callback replaces arrays without mutating its input leaves. That does not make every possible tree-map callback free of side effects.
  • Do not assume one transformation is valid for every NNX model. Check variable roles, leaf dtypes, paths and any shared references first.
  • If JIT-compiling this setup function, mark seed and features static because their integer values control initialization and shapes.