[Feature] Add inplace=True to pad - #1706
Merged
Merged
Conversation
`pad(td, ...)` currently allocates a fresh tensordict alongside the input, so peak memory is ~2x the leaves until the caller's reference to the original drops. With `inplace=True` the function rebinds each leaf inside the input as it goes, releasing the old storage before the next leaf's pad runs, and updates the batch_size at the end. Peak overhead collapses to roughly one leaf's worth. LazyStackedTensorDict can't have its batch_size mutated directly, so the inplace path pads each constituent along non-stack dims in place and grows the stack dim by appending/prepending zero-filled copies of the edge constituents. Also adds `td.pad(...)` as a method on TensorDictBase, registers it on tensorclass, and documents the project's `method_` (same-storage) vs `inplace=True` (same-container) convention in CLAUDE.md and docs/source/overview.rst.
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
Feature
`pad(td, ..., inplace=True)` mutates the input leaf by leaf. If a later leaf's pad raises, earlier leaves are already padded and the batch_size still reflects the old shape, leaving the tensordict inconsistent. A full transactional rollback would require keeping every old leaf alive until the whole pass succeeds, which would defeat the 1x memory contract. This adds a `safe` kwarg (default `True`) that runs a non-mutating pre-flight walk over every leaf, validating the per-leaf preconditions of torch.nn.functional.pad (sufficient dims, no negative output size) and recursing into nested tensordicts and lazy-stack constituents. Realistic user errors (bad pad widths, unpaddable leaves) now raise before any mutation. `safe=False` skips the walk for a small speedup when the inputs are known to be valid; the docstring warns that OOM-style failures mid-loop still corrupt the tensordict regardless of `safe`.
test_tensorclass_stub_methods enforces that every public method on TensorDict has a matching signature in tensorclass.pyi. Add the pad stub. Bundles a pure-formatting pass from the linter.
4 tasks
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
1 participant
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
inplace=Truetotensordict.padso it pads each leaf in-place inside the input tensordict and releases each old storage as soon as the replacement is written. Peak memory fortd = pad(td, ...)-style usage drops from ~2x the leaves to ~1x.padas a method onTensorDictBase(and registers it ontensorclass).LazyStackedTensorDictis handled by padding each constituent along the non-stack dims in place and growing the stack dim by appending/prepending zero-filled copies of the edge constituents — so the lazy stack's identity is preserved even though itsbatch_sizesetter forbids resizing.method_(same-storage) vsinplace=True(same-container) convention inCLAUDE.mdanddocs/source/overview.rst, and the tensorclass dispatch list registerspadso the fallback warning is no longer triggered.The motivation:
pad(td, ...)currently builds a fresh tensordict viaTensorDict._new_unsafe({}, new_batch_size, ...)and populates it leaf-by-leaf, holding the originaltdalive for the duration of the call. Common usage istd = pad(td, ...)where the caller wants the original released, but the function-local reference keeps it alive until return — peak ~2x. Withinplace=Truethe rebind happens inside the input, so each old leaf is freed before the next leaf's pad allocates.Test plan
pytest test/test_tensordict.py -k pad— 58 passed, 1 unrelated skippytest test/test_tensordict.py::TestGeneric— 279 passedtd.pad([0, 0, 0, 1], inplace=True) is tdreturnsTrue.Follow-ups (not in this PR)
Other shape-changing ops with the same 2x-memory profile — candidates for the same
inplace=Truetreatment in future PRs:gather(already hasout=),repeat/repeat_interleave,roll,reshape/flatten/unflatten(only when the underlying call materializes a copy),contiguous,where. None warrant an underscore variant, since they all change shape/dtype/layout.