[Performance] Add opt-in fast_stack for TensorDict

[vmoens]

Summary

Adds tensordict.fast_stack (and TensorDictBase.fast_stack classmethod), a strict, opt-in variant of stack that walks all input TensorDicts in lockstep instead of doing per-leaf string lookups across each input. Same shape as the update_at_ → copy_at_ split: both functions remain available, and fast_stack raises RuntimeError when its preconditions aren't met so callers can decide between the two explicitly.

tensordict.stack is unchanged. The fast path is purely additive.

Speedup

Benchmark: stacking N homogeneous TensorDicts with M leaves at depth levels. Mean per call (lower is better):

Shape (N, M, depth)	stack	fast_stack	Speedup
(8, 20, 0)	127 µs	32 µs	4.0x
(64, 20, 0)	869 µs	151 µs	5.7x
(8, 10, 2)	207 µs	34 µs	6.1x
(64, 10, 2)	1,340 µs	151 µs	8.9x

The win comes from skipping the per-leaf `_get_str` calls and per-leaf Python shape validation in `_stack`. Speedup grows with N (number of TDs) and depth (nesting).

What fast_stack accepts

• Plain TensorDict at root, with arbitrary nesting.
• Tensorclass at root or at any nested leaf (including tensorclass-wrapping-LazyStackedTensorDict).
• LazyStackedTensorDict at root or nested, when all inputs agree on stack_dim and inner-TD count.
• NonTensorData / NonTensorStack / MetaData (via _pass_through_cls).
• nn.Parameter, MemoryMappedTensor, UninitializedParameter, UninitializedBuffer, UnbatchedTensor leaves.
• Differing key insertion orders across inputs (lockstep zip when orders match, td0-driven dict lookup mid-traversal on mismatch).

What fast_stack raises on (use stack instead)

• Mixed root types (e.g. TensorDict + LazyStackedTensorDict).
• Mismatched batch_size, device, or key set.
• LazyStackedTensorDict with mismatched stack_dim or inner-TD count.
• PersistentTensorDict, _SubTensorDict, TensorDictParams, TypedTensorDict, or any other TensorDict subclass.
• API kwargs not implemented: out=, device=, strict, contiguous, maybe_dense_stack.
• lazy_legacy=True global mode is ignored — fast_stack always returns dense.

Commit history

The branch is split into 8 logical commits for review. The first lands a strict baseline; the next six progressively loosen support (key ordering, tensor subclasses, tensorclass + pass-through, lazy stack, hot-path inlining, tensorclass-wrapping-lazy); the last cleans up the error message and adds negative tests.

Test plan

• pytest test/test_tensordict.py -k fast_stack — 30 new tests pass (1 h5py-skipped). Covers equivalence vs torch.stack on flat/nested TDs, tensorclass at root and nested, tensorclass-wrapping-lazy, lazy stack at root and nested, NonTensorData at root and nested, nn.Parameter / MemoryMappedTensor / UninitializedParameter / UnbatchedTensor leaves, key-order mismatch (positive), and rejections for mixed root types / PersistentTensorDict / device mismatch / batch_size mismatch / key set mismatch / lazy stack_dim mismatch / dim out of range / empty input.
• pytest test/test_tensordict.py -k "stack or Stack" — 1684 stack-related tests pass total; existing stack paths unchanged.
• torch.compile(fullgraph=True) works on fast_stack over nested TDs.
• Benchmark via benchmarks/common/common_ops_test.py::test_fast_stack_homogeneous parametrized over (num_tds, num_keys, depth).

Stacked on

This PR is stacked on #1692. Once #1692 merges this should auto-rebase onto main.

🤖 Generated with Claude Code