[BugFix] Keep recursive TensorDict metadata consistent - #1731
Merged
Conversation
vmoens
force-pushed
the
codex/recursive-properties-checker
branch
from
June 30, 2026 11:19
ff58b8f to
3a7c2cd
Compare
vmoens
force-pushed
the
codex/recursive-properties-checker
branch
from
June 30, 2026 11:42
89c98f1 to
b4db4cc
Compare
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.
Description
_check_recursive_propertiesTensorDict utility that recursively validates node and leaf metadata, with aggregated error messages and opt-outs for batch size, shape, device, lock, names, non-tensor, and unbatched checks.UnbatchedTensortransformation tests, allTestTensorDictsBasefixture factories, a broadTestTensorDictscommon-op sweep, and a tensorclass common-op sweep.UnbatchedTensorTensorDict-facingbatch_sizemetadata synchronized while preserving payload storage across construction/validation, stack/cat, lazy stack, indexing, unbind/split/chunk, gather, masked select, view/reshape/unflatten, transpose/permute/movedim, and lazy custom-op views.Motivation and Context
TensorDict transformations can preserve an
UnbatchedTensorpayload while changing the parent TensorDict batch shape. The payload should remain shape-independent, but the wrapper metadata exposed through TensorDict should match the current parent context so downstream checks do not observe stale batch sizes.The checker coverage added here is deliberately broader than the patched class itself: it now runs through the centralized
TestTensorDictsBasefixture factories plus common creation/transformation/modification paths for TensorDicts and tensorclasses. Under that broader coverage, the only real implementation bugs found in this pass wereUnbatchedTensorTensorDict-facing metadata issues; I did not find additional class-specific invariant violations in dense TensorDict, lazy stack, persistent/memmap/HDF5-backed TensorDicts, or tensorclasses.Concrete issues patched while adding the checker:
UnbatchedTensorvalues inserted into a TensorDict kept empty/stale metadata instead of adopting the parent batch size.UnbatchedTensorleaves returned stale metadata, especially when source leaves had empty metadata.getitem,unbind,split,chunk,gather,masked_select,view/reshape,unflatten,transpose/permute/movedim, and legacy lazy custom-op views) preserved payload storage but did not refresh TensorDict-facingbatch_sizemetadata.No public API is added; the checker is private for test/debug rollout.
Runtime impact
Local measurements on this branch:
6.0 usper checker call; all representative added checker calls together: median0.11 msper pass.TestTensorDictsBasefixture checks measured locally at roughly6.7-26.5 usper fixture construction across representative fixture types.test/tensordictshard before adding centralized fixture checks:7790 passed, 898 skipped in 13.69s.test/tensordictshard after adding centralized fixture checks:7790 passed, 898 skipped in 14.35s(real 14.94s).test/tensordictshard after the broader common-op coverage:7809 passed, 898 skipped in 13.94s(real 14.47s). This is within local run-to-run noise; expected overhead for the full shard is under about one second.8 passed, 24030 deselected in 1.12s; the added tensorclass common-op coverage is included intest/tensorclass/test_tensorclass.py(148 passed, 1 skipped in 2.74s).Types of changes
Checklist
Tests
PYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/tensordict -qPYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/utils/test_utils.py -k 'check_recursive_properties' -qPYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/tensorclass/test_tensorclass.py -quvx pre-commit run --files test/_utils_internal.py tensordict/base.py tensordict/_td.py tensordict/_torch_func.py tensordict/_lazy.py tensordict/utils.py test/tensordict/test_nontensor.py test/tensordict/test_methods.py test/utils/test_utils.py test/tensorclass/test_tensorclass.py