Skip to content

← Flax NNX step 15 of 100

Medium Primitives

NNX Pure State Update

Jump to the task

Why this matters

You can transform a model's state without changing the original model. This is useful when comparing modified weights, implementing an explicit update rule or passing state through a function.

Here, split a Linear layer into structure and state, multiply its array leaves by a scale, and reconstruct a layer from the transformed state. The original values stay unchanged because the mapping callback computes new arrays instead of assigning into the original Variables.

This is one way to work with NNX. Object-aware transformations and optimizers also support model objects and state updates; they do not all require you to write a manual split/map/merge sequence.

Worked example

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

model = nnx.Linear(3, 2, rngs=nnx.Rngs(0))
model.bias[...] = jnp.array([1.0, -1.0])
original_kernel = model.kernel[...]
original_bias = model.bias[...]
graphdef, state = nnx.split(model)

new_state = jax.tree.map(lambda leaf: leaf * 2.0, state)
merged = nnx.merge(graphdef, new_state)

assert jnp.array_equal(model.kernel[...], original_kernel)
assert jnp.array_equal(model.bias[...], original_bias)
assert jnp.array_equal(merged.kernel[...], 2.0 * original_kernel)
assert jnp.array_equal(merged.bias[...], jnp.array([2.0, -2.0]))

The example uses a nonzero bias so you can see that both arrays are scaled. The exercise below uses Linear's default zero bias.

Tree mapping reconstructs the pytree structure around the transformed leaves, including the Variable wrappers. Merge then reconstructs the model. Purity comes from the operations performed, not from calling a variable new_state: a callback that mutates a captured object would still have a side effect. Assigning an existing state to a new Python name is not a copy.

Which state should be scaled?

All state in this Linear layer consists of parameters. A larger model may also contain counters, random keys or statistics, which should not all receive the same arithmetic update.

For a parameter-only change, select nnx.state(model, nnx.Param) first. Mapping that subset and calling nnx.update(model, changed_params) is another useful pattern, but it intentionally updates the existing model. An actual optimizer also needs gradients and may maintain its own state; this exercise does not implement a complete training step.

Problem

Implement pure_state_update(seed, x, features, scale):

  1. Construct nnx.Linear(x.shape[-1], int(features), rngs=nnx.Rngs(int(seed))) with its default initializer and bias configuration.
  2. Split the model into graph definition and state.
  3. Multiply every state array leaf by scale using a tree mapping.
  4. Merge the transformed state and return the reconstructed model's output on x.

Input x is a nonempty finite float32 vector. Features is a positive integer-valued scalar and seed is a nonnegative integer-valued scalar; both may arrive as floats. Scale is a finite scalar. Return a float32 vector of length features.

The calculation is x @ (scale * kernel) + scale * bias. For these finite parameters, scale zero produces zeros, scale one preserves the original calculation, and a negative scale reverses its sign. Algebraic rearrangements can differ slightly through floating-point rounding.

Common pitfalls

  • Use multiplication for scaling; adding scale implements a different update.
  • Keep scale as a JAX value. Python float conversion breaks tracing and differentiation when scale is dynamic.
  • Do not mutate a Variable in the original state while building the new state.
  • A transformed state is not callable. Merge it before applying the layer.
  • If JIT-compiling the setup function, mark seed and features static. Scale can remain a dynamic argument.

Numerical outputs cannot establish that you used state mapping or preserved the original model. Run the example's before/after assertions to verify the intended state behavior.