[KennethGrace]

Closes #20.

GatedDeltaNetReadCache and the kernel's ReadCache came after #12 registered the state containers, so torch still treated a cache as one opaque leaf. That never showed while a cache was made and used inside one vmapped call, or shared unbatched across sets. It did show once the write and the re-read happened in separate batched calls, which is what stack_module_state + vmap(functional_call(...)) needs. G writers could not return their caches:

ValueError: vmap(write, ...): `write` must only return Tensors, got type <class 'lumen.gdn.layer.GatedDeltaNetReadCache'>

and so G re-readers had nothing to take back in at in_dims=0. Both now work.

The change

ReadCache.seq_len is the UndertowState.seen case again. It is a permanent int that drives control flow: where a re-read trims its padded output, and what a query's length is checked against. register_dataclass would make it a leaf, and vmap refuses it. So it goes in the treespec context, and it gets no default, for the reason lumen/pytree.py already gives for seen.

One helper rather than a second copy. The Undertow registration's hand-rolled flatten is now _register_with_context(cls, context). It takes the context field by name and reads the children off the dataclass, so the keyed and unkeyed flattens come from one field list and cannot disagree. A field added to either class later becomes a child automatically. If the children were listed by hand, a new field would be dropped by the flatten and silently rebuilt from its default. Now the worst a new non-tensor field can do is be refused by vmap, which is loud. UndertowState goes through the same helper, and its flatten is unchanged: children [keys, values], context seen.

The outer GatedDeltaNetReadCache needs nothing special. Its trajectory flattens through its own registration, and its conv is None on a pass that began a sequence, the same None-or-tuple case register_dataclass already handles for GatedDeltaNetState.conv.

REGISTERED grows from five names to seven. The caches are appended, so the five states keep their positions.

Numerically inert, checked rather than assumed

Weights, a vmapped two-chunk stream through an Undertow-plus-GDN trunk (which goes through the rewritten Undertow flatten), and a GDN re-read all hash identically on main and on this branch. Registration changes a registry inside torch. No tensor changes value, and no checkpoint key moves.

Tests

6 new, CPU-only, each run for a first pass (conv is None) and a continued one (conv is a tuple):

• Round trip. seq_len survives flatten/unflatten, every leaf is a tensor, and the rebuilt cache re-reads the original output exactly (atol=0).
• Paths. tree_flatten_with_path reaches every leaf through both registrations (.trajectory.unfolded … .conv[2]), and a first pass has no .conv leaf.
• The capability. G vmapped writes, then G vmapped re-reads in a separate call, match the layers run one at a time. Every cache leaf comes back with the G axis, conv included.

Against main's pytree.py, with the skip guard disabled, all 7 new or changed tests fail with the errors the issue reports. The 22 existing tests pass both ways.

On exactness, measured on one machine (CPU, torch 2.9.1): re-reading under vmap over caches written one at a time is exact, max abs difference 0.0, which matches the issue's probe. With the writes batched too, the difference is about 1.8e-7. That is fp32 rounding from batching the chunkwise scan, the same thing the existing set-by-set test tolerates. The new test uses that test's tolerance (1e-5) rather than asserting zero.

Degradation, checked by stubbing torch's entry points before import:

• With the private register_pytree_node raising, REGISTERED drops to five and the layer writes and re-reads exactly as before. As with a BlockState holding an unregistered UndertowState, the outer cache then flattens to one opaque leaf, and REGISTERED shows ReadCache as missing.
• With the keyed variant refused, registration still completes, and only path traversal is lost.

591 passing with the CUDA tests deselected. The GPU-marked tests were not run for this change: registration is pure Python over tensor references and never looks at a device.

Docs

• docs/PYTREE.md gains a read-cache section with a working write-then-re-read example, and the REGISTERED listing is updated.
• docs/design/GATED_DELTANET.md §3.9 said the cache was "not registered … and is its own change". It now describes the registration.
• README module table updated.

Release

0.8.1, in the last commit. Additive and numerically inert, so a patch release.

🤖 Generated with Claude Code