nnx.capture fails with nnx.scan when the model contains nnx.List (mixed str/int pytree dict keys)
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
### 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
```python
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.
Contributor guide
Assessment
This issue has not been assessed yet.