[Feature] Add inplace=True to gather / repeat / repeat_interleave / roll / reshape / flatten / unflatten / contiguous - #1707
Merged
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #1706 (which added
inplace=Truetopad). 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:gather,repeat,repeat_interleave,roll.reshape,flatten,unflatten,contiguous.Each op gains an
inplace: bool = Falsekwarg. WhenTrue, the tensordict's object identity and key set are preserved, leaves are replaced one at a time insideself, and the old leaf storages are released as their replacements are written — so the commontd = td.op(...)pattern drops from ~2x peak memory to ~1x.A new shared helper
TensorDictBase._inplace_rebind_leavesencapsulates the iteration / rebind /_change_batch_sizesequence so each op only provides the per-leaf transform, the nested-TD recursion, and the newbatch_size.Design notes
NotImplementedErroroninplace=Truefor 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.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 returnselffromTensor.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): requiresindex.ndim == self.batch_dimsso the result keeps the same number of batch dims. Mutually exclusive without=.repeat_interleave(inplace=True): requires an explicitdimandndim > 0because thedim=None/ scalar paths internally reshape, which would changendim.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