fix(util): traverse container subclasses so ModelOutput values are not silently skipped

[Hotragn]

Summary

util.apply and util.applyn decide which values nnsight is allowed to touch. They matched the builtin containers exactly:

data_type = type(data)

if data_type == list: ...
elif data_type == tuple: ...
elif data_type == dict: ...

so every subclass of them fell through to return data untouched. A HuggingFace ModelOutput is an OrderedDict subclass, and it is what every top-level transformers module returns — which means model.output and model.transformer.output were invisible to all five callers of these helpers. namedtuple, OrderedDict, defaultdict and torch.return_types.* module outputs were skipped for the same reason.

Nothing raised. The traversal simply found no tensors, so each caller degraded silently.

The five symptoms

Reproduced on gpt2, CPU, dev @ efe4244:

Caller	Symptom
Batcher.narrow	An invoke reading model.output in a batched trace got the whole batch, and both invokes got the same object (a.logits is b.logits → True)
Batcher.swap	Writing through model.transformer.output inside one invoke landed on every other invoke's rows
hooks.cache_output_hook / cache_input_hook	tracer.cache(...) stored the full batch for a ModelOutput module
Cache.add	cache(dtype=...) / cache(device=...) left those tensors unconverted; the default detach=True never fired, so the cache pinned the autograd graph
nnsight_forward's __nnsight_skip__ merge	A multi-invoke tracer.skip() returned one invoke's value instead of the concatenation

Read leak

with m.trace() as tracer:
    with tracer.invoke(A):
        a_h0    = m.transformer.h[0].output[0].save()   # plain tuple
        a_inner = m.transformer.output.save()           # BaseModelOutputWithPast
        a_top   = m.output.save()                       # CausalLMOutputWithPast
    with tracer.invoke(B):
        ...

h[0].output[0]            (plain tuple)   got=(1, 10, 768)      expected=(1, 10, 768)      OK
transformer.output.last_hidden_state      got=(2, 10, 768)      expected=(1, 10, 768)      *** WRONG ***
model.output.logits                       got=(2, 10, 50257)    expected=(1, 10, 50257)    *** WRONG ***

Write leak

Zeroing hidden states in invoke 0 only:

invoke0 logits changed?  True   (expected True)
invoke1 logits changed?  True   (expected False)   <- clobbered

Cache

cache(dtype=torch.float16):
  model.transformer        BaseModelOutputWithPast   dtypes=['torch.float32']  *** NOT CONVERTED ***
  model.transformer.h.0    tuple                     dtypes=['torch.float16']  OK
  model.lm_head            Tensor                    dtypes=['torch.float16']  OK

cache in a 2-invoke trace:
  model.transformer        BaseModelOutputWithPast   (2, 10, 768)  *** batch dim not 1 ***

This is the failure mode that worries me most: a batched-invoke user reading model.output gets a plausible tensor with a silently wrong batch dimension, and cache(device=...) silently leaves activations on the GPU it was asked to move them off.

Fix

Recurse into list / tuple / dict subclasses too, rebuilding the concrete type so callers still get a ModelOutput back rather than a plain dict.

• Mappings and lists are rebuilt by mutating a shallow copy, not by reconstructing the type from its items — subclasses like ModelOutput are dataclasses whose __init__ does not accept a mapping. Going through __setitem__ also preserves ModelOutput's item/attribute mirroring, so .logits stays in sync with ["logits"] (asserted in the tests).
• Tuples use _make for namedtuples and the iterable constructor otherwise (torch.Size, torch.return_types.* structsequences), falling back to a plain tuple only when the subclass genuinely cannot be rebuilt.

The exact-type fast paths are untouched, so the hot path is unchanged: the new branches only run for values that previously fell straight through to return data. Batcher.narrow also still early-returns before any of this when needs_batching is False, so single-invoke traces do no extra work at all.

Deliberately out of scope: collections.UserDict / UserList and transformers.BatchEncoding are not dict/list subclasses (they are MutableMapping/MutableSequence), and objects like DynamicCache are not containers at all. Those keep today's behaviour — widening to the abc types is a bigger semantic question (str is a Sequence) and none of the symptoms above need it.

Verification

Each invoke's narrowed ModelOutput now matches the same prompt traced alone to a relative error of 2e-7, while differing from the other invoke's reference by 1.7e-1 — so the slice is right, not merely the shape.

tests/test_container_types.py adds 33 tests: the traversal itself (apply/applyn over 9 container shapes, type preservation, inplace, torch.Size, structsequences, a tuple subclass with an incompatible constructor, non-containers, strings) plus the five behaviours above. 18 of the 33 fail on dev and all 33 pass with this change.

Full CI test list on CPU: 1 failed, 276 passed. The one failure is test_lm.py::TestGradients::test_backward_with_multiple_invokers, which is already red on dev and main and is unrelated to this change (addressed separately in #693).

_{Run with transformers 5.15.1 / torch 2.9.1 on CPU. Verified this change introduces no failure difference under transformers 4.57.3 either.}

[Hotragn]

Closing this — checked it against 0.8 per your note in #696, and it's already fixed there.

util.apply on 0.8 is a full rewrite: isinstance instead of type(data) ==, exact-type rebuild for dict subclasses via .copy(), _fields for namedtuples, plus set / frozenset / bounded object descent. applyn is gone entirely. Verified directly against 0.8:

BaseModelOutputWithPast    type=BaseModelOutputWithPast  transformed=True
CausalLMOutputWithPast     type=CausalLMOutputWithPast   transformed=True
OrderedDict / defaultdict / namedtuple / return_types    transformed=True
.last_hidden_state is ["last_hidden_state"]: True
original untouched: True

And the behaviour I was fixing holds on 0.8 — batched invokes reading model.output:

solo batch=1  invoke0 batch=1  invoke1 batch=1
invoke0 is invoke1 object: False
invoke0 vs solo rel_err: 2.132e-07

tests/test_util.py already covers the traversal and tests/test_batching.py::test_batched_matches_solo_logits already covers the read path through a ModelOutput, so there's nothing left here worth porting. Nothing for you to action.