Skip to content

[Feature] Add inplace=True to gather / repeat / repeat_interleave / roll / reshape / flatten / unflatten / contiguous - #1707

Merged
vmoens merged 2 commits into
pytorch:mainfrom
vmoens:inplace-shape-ops
May 19, 2026
Merged

[Feature] Add inplace=True to gather / repeat / repeat_interleave / roll / reshape / flatten / unflatten / contiguous#1707
vmoens merged 2 commits into
pytorch:mainfrom
vmoens:inplace-shape-ops

Conversation

@vmoens

@vmoens vmoens commented May 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

Follow-up to #1706 (which added inplace=True to pad). Extends the same leaf-by-leaf-rebind pattern to the rest of the shape-changing ops that were flagged in that PR's punch list:

  • Always-copy ops: gather, repeat, repeat_interleave, roll.
  • Sometimes-view ops: reshape, flatten, unflatten, contiguous.

Each op gains an inplace: bool = False kwarg. When True, the tensordict's object identity and key set are preserved, leaves are replaced one at a time inside self, and the old leaf storages are released as their replacements are written — so the common td = td.op(...) pattern drops from ~2x peak memory to ~1x.

A new shared helper TensorDictBase._inplace_rebind_leaves encapsulates the iteration / rebind / _change_batch_size sequence so each op only provides the per-leaf transform, the nested-TD recursion, and the new batch_size.

Design notes

  • LazyStackedTensorDict raises NotImplementedError on inplace=True for these ops; its overrides reshape/flatten/unflatten/repeat into the stack structure itself, which can't be safely rebound in place. The error message suggests .to_tensordict() as the escape hatch.
  • View branch (reshape, flatten): Tensor.<op> may return a view when memory layout permits. The rebind still happens (object identity + batch_size are correctly updated), but the new leaf shares storage with the original so the memory benefit does not materialize. The docstrings call this out.
  • contiguous(inplace=True): leaves that are already contiguous return self from Tensor.contiguous(), so the rebind is a no-op for those keys. Useful when the leaves are a mix of contiguous and non-contiguous tensors.
  • gather(inplace=True): requires index.ndim == self.batch_dims so the result keeps the same number of batch dims. Mutually exclusive with out=.
  • repeat_interleave(inplace=True): requires an explicit dim and ndim > 0 because the dim=None / scalar paths internally reshape, which would change ndim.
  • No safe= kwarg here (unlike pad). For these ops there's no realistic per-leaf pre-flight check that doesn't require a CUDA sync — the only mid-loop failure mode is OOM, which can't be pre-validated.

Test plan

  • `pytest test/test_tensordict.py -k "repeat or roll or gather or reshape or flatten or unflatten or contiguous or inplace"` — 672 passed, 26 unrelated skips
  • `pytest test/test_tensordict.py::TestGeneric -k inplace` — 33 new tests
  • `pytest test/test_tensorclass.py` — 144 passed (stub tests pass with new `inplace` kwargs added to `tensorclass.pyi`)
  • New tests cover, per op: `inplace=True` returns `self`; leaf-for-leaf equivalence with the `inplace=False` path; `LazyStackedTensorDict` raises `NotImplementedError`; for `repeat` and `roll`, a weakref-based check that old leaf storages are released. `gather` also checks the `out=` + `inplace=True` conflict and the ndim-mismatch error.

…oll / reshape / flatten / unflatten / contiguous

Extends the same leaf-by-leaf-rebind pattern that pad already uses to
the other shape-changing ops. Each op gains an ``inplace`` kwarg
(default ``False``); when True the tensordict's object identity and
key set are preserved, leaves are replaced one at a time, and the
old leaf storages are released as their replacements are written.
Peak memory for the common ``td = td.op(...)`` pattern drops from
~2x the leaves to ~1x.

A shared ``TensorDictBase._inplace_rebind_leaves`` helper encapsulates
the iteration / rebind / ``_change_batch_size`` sequence so each op
only has to provide the leaf transform, the nested-TD recursion, and
the new batch_size.

LazyStackedTensorDict raises NotImplementedError on inplace=True for
these ops (its overrides reshape/flatten/unflatten/repeat into the
stack structure, which can't be safely mutated in place). The error
message points at ``.to_tensordict()`` as the escape hatch.

For reshape/flatten the underlying ``Tensor.<op>`` may return a view
when memory layout permits; in that case the rebind happens but the
new leaf shares storage with the original, so the memory benefit
does not materialize. The docstrings call this out. ``contiguous``
is similar: leaves that are already contiguous return ``self`` and
the rebind is a no-op.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 18, 2026
@github-actions github-actions Bot added Feature New feature Test tensorclass and removed Feature New feature labels May 18, 2026
@github-actions github-actions Bot added the Feature New feature label May 19, 2026
@vmoens
vmoens merged commit a058f4a into pytorch:main May 19, 2026
26 of 42 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Feature New feature tensorclass Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant