google / google/flax

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

Open
#5,571 1 comment 0 reactions 1 assignee Claimed by @samanklesaria View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.