Mar 4, 2026

Cristian Garcia, Samuel Anklesaria, Flax Team

check_no_aliases(args=args)

The new simplified NNX transform implementations don’t support this argument. However, to keep supporting the behavior, a new `nnx.transform_metadata` transform is introduced that can be inserted to get back the same results. TODO: mention it works on `jax.vmap`.

Previously, `Module.sow` used graph updates to capture intermediate values during computations and propagate them outside, it was used in conjunction with `nnx.pop` to log and extract intermediates:

class Foo(nnx.Module):

def __call__(self, x):

self.sow(nnx.Intermediate, "y_mean", jnp.mean(x))

return x

model = Foo()

result = model(x)

intermediates = nnx.pop(model, nnx.Intermediate) # extract intermediate values

To achieve the same without graph updates we’ve added a new `nnx.capture` API which allows for a similar workflow.

class Foo(nnx.Module):

def __call__(self, x):

self.sow(nnx.Intermediate, "y_mean", jnp.mean(x))

return x

model = Foo()

result, intermediates = nnx.capture(model, nnx.Intermediate)(x)

In general, `nnx.capture` takes a function or Module to be transformed, a `nnx.Variable` subclass to collect, and an optional `init` argument to initialize the collected state, which will be stored within `nnx.Variable` objects. `nnx.capture` creates a `__captures__: tuple[Variable, ...]` attribute on each `Module` instance, each Variable in `__captures__` contains a dictionary which `sow` and `perturb` populate.

Similarly, `Module.perturb` was previously used to extract the gradients of intermediate values. This was done in two steps: initializing a perturbation state by running a module once, and then passing the perturbation state as a differentiable target to `grad`.

class Model(nnx.Module):

def __call__(self, x):

x = self.perturb('grad_of_x', x)

...

return y

def train_step(model, optimizer, x, y):

model(x) # Initialize perturbation state

def loss_fn(model):

y_pred = model(x)

return jnp.mean((y_pred - y) ** 2)

diff_state = nnx.DiffState(0, (nnx.Param, nnx.Perturbation))

grads = nnx.grad(loss_fn, argnums=diff_state)(model)

grads, interm_grads = nnx.state(grads, nnx.Param, nnx.Perturbation)

optimizer.update(model, grads)

nnx.pop(model, nnx.Perturbation) # clean up perturbations

return interm_grads

Similar pattern can be used with `nnx.capture` during both perturbation initialization and when running the forward pass to insert the differentiable perturbations state. In this version explicitly pass the `perturbs` state as a separate argument and use `argnums` to specify that both arguments are differentiable:

def train_step(model, optimizer, x, y):

_, perturbs = nnx.capture(model, nnx.Perturbation)(x) # init perturbations

def loss_fn(model, perturbs):

y_pred = nnx.capture(model, init=perturbs)(x)

return jnp.mean((y_pred - y) ** 2)

grads, interm_grads = nnx.grad(loss_fn, argnums=(0, 1))(model, perturbs)

optimizer.update(model, grads)

return interm_grads