nnx.capture fails with nnx.scan when the model contains nnx.List (mixed str/int pytree dict keys)

[maxencefaldor]

System information

• Flax 0.12.7, JAX 0.10.1, Python 3.14, macOS arm64 (CPU)

Problem

nnx.capture fails when the wrapped function runs an nnx.scan and the model contains an nnx.List (directly or via nnx.Sequential):

ValueError: Comparator raised exception while sorting pytree dictionary keys.
...
TypeError: '<' not supported between instances of 'str' and 'int'

The same model captures fine eagerly and under nnx.jit alone; the same scan captures fine when the model stores its layers as plain attributes instead of an nnx.List. Only the combination capture + scan + nnx.List fails.

Reproduction

import jax.numpy as jnp
from flax import nnx


class Model(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.layers = nnx.List([nnx.Linear(4, 4, rngs=rngs) for _ in range(2)])

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        self.sow(nnx.Intermediate, "out", x)
        return x


def rollout(model, x):
    state_axes = nnx.StateAxes({nnx.Intermediate: 0, ...: nnx.Carry})
    return nnx.scan(
        lambda m, x: m(x), in_axes=(state_axes, nnx.Carry), out_axes=nnx.Carry, length=3
    )(model, x)


x = jnp.ones(4)

# eager: OK
nnx.capture(Model(nnx.Rngs(0)), nnx.Intermediate)(x)

# nnx.jit: OK
nnx.capture(nnx.jit(lambda m, x: m(x)), nnx.Intermediate)(Model(nnx.Rngs(0)), x)

# nnx.scan: fails
nnx.capture(rollout, nnx.Intermediate)(Model(nnx.Rngs(0)), x)

Output:

capture(model)               OK
capture(jit(model))          OK
capture(scan rollout)        ValueError: Comparator raised exception while sorting pytree
                             dictionary keys.  (cause: TypeError: '<' not supported between
                             instances of 'str' and 'int')

Replacing the nnx.List with two plain attributes (self.l0, self.l1) makes the scan case pass and the sown value comes out stacked as expected, so the scan/StateAxes usage itself appears correct.

Analysis

capture()'s wrapper sets m.__captures__ = pytreelib.data(...) on every module yielded by iter_modules(module) (flax/nnx/module.py, wrapper). iter_modules yields nnx.List (it is a Module subclass), so the List node — whose existing children are keyed by integers — gains a string-keyed entry. When jax.lax.scan later flattens the carry (first hit in api_util.debug_info → flatten_with_path), the dict with mixed int/str keys cannot be key-sorted and the comparator raises.

A possible fix is to skip integer-keyed container modules (nnx.List/nnx.Dict-style nodes) when planting __captures__ — a sow on such a container has no stable attribute name anyway — or to store the capture buffers out-of-band rather than as attributes.

What you expected to happen

The scan case behaves like the eager and jit cases: (result, intermediates) with the sown value stacked along the scan axis.