Skip to content

← Flax NNX step 12 of 100

Medium Primitives

NNX State Split/Merge

Jump to the task

Why this matters

nnx.split(model) separates a model's structure from its state. The graph definition describes how to reconstruct the model; the state holds its changing values in a JAX pytree. nnx.merge(graphdef, state) puts the two back together.

This is useful when you want to pass state explicitly through a JAX transformation. You can also inspect or transform the state separately from the module that uses it.

This exercise checks a round trip for a Linear layer. Its forward call reads the kernel and bias without changing them, so reconstructing it from unchanged state preserves its output on the same input.

Worked example

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

model = nnx.Linear(3, 4, rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
original_output = model(x)

graphdef, state = nnx.split(model)
merged = nnx.merge(graphdef, state)
merged_output = merged(x)

assert merged is not model
assert jnp.array_equal(original_output, merged_output)
assert jnp.max(jnp.abs(original_output - merged_output)) == 0

The new module has the same layer configuration and parameter values. Merge does not call Linear's initializer to draw a new random kernel. This is a reconstruction operation, not a general-purpose deep copy of arbitrary Python objects or external resources.

Passing state through JIT

Continuing the example, define the compiled function once and pass state as an argument:

@jax.jit
def apply_state(current_state, inputs):
    reconstructed = nnx.merge(graphdef, current_state)
    return reconstructed(inputs)

compiled_output = apply_state(state, x)
assert jnp.allclose(compiled_output, original_output, rtol=1e-5, atol=1e-6)

This is a small functional wrapper for this read-only Linear call. It is not an implementation of nnx.jit. A wrapper for a state-changing call must also return the updated state if the caller needs to retain those changes. The graph definition is captured by the function; state is an explicit argument, so a later call can receive updated parameter values.

Problem

Implement split_merge_round_trip(seed, x, features):

  1. Construct nnx.Linear(x.shape[-1], int(features), rngs=nnx.Rngs(int(seed))) using its default initializer and bias configuration.
  2. Compute original_output = model(x).
  3. Split the model, then merge the returned graph definition and state.
  4. Compute the reconstructed model's output on the same x.
  5. Compute the largest absolute elementwise difference between the outputs.
  6. Return a float32 array containing the original output sum, reconstructed output sum and largest difference, in that order.

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. The return shape is (3,). For the supplied inputs, both sums agree and the difference is zero.

Keep output values as JAX arrays and combine them with jnp.stack. Converting them to Python floats prevents tracing under JIT.

The numerical fixtures check the sums and difference. They cannot prove that your code used split and merge, so also run the example and inspect the reconstructed model while practicing those operations.

Common pitfalls

  • Keep the graph definition compatible with the state. State may contain nested submodule entries; it is not necessarily a flat dictionary.
  • Do not assume every model repeats its previous output. Dropout, counters or other state-changing calls need careful control of state and randomness. This exercise uses a deterministic Linear layer with unchanged parameters.
  • Supply the state values you intend to use. A separately saved old snapshot can reconstruct old parameters; caching state is not automatically safe.
  • Default merge can share Variable objects with the supplied state and the original model. A new module is not necessarily an independent model. Replacing a shared kernel value can affect both modules. Use nnx.merge(graphdef, state, copy=True) when you need copied Variables; this is still not a general deep-copy guarantee for arbitrary objects.
  • Do not recreate a decorated JIT function on every call. Define it once for the graph structure you intend to reuse.
  • If JIT-compiling the exercise's setup function, mark seed and features static because their integer values control initialization and shapes.