Skip to content

[BugFix] Keep recursive TensorDict metadata consistent - #1731

Merged
vmoens merged 5 commits into
pytorch:mainfrom
vmoens:codex/recursive-properties-checker
Jun 30, 2026
Merged

[BugFix] Keep recursive TensorDict metadata consistent#1731
vmoens merged 5 commits into
pytorch:mainfrom
vmoens:codex/recursive-properties-checker

Conversation

@vmoens

@vmoens vmoens commented Jun 30, 2026

Copy link
Copy Markdown
Collaborator

Description

  • Add a private _check_recursive_properties TensorDict 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.
  • Exercise the checker in focused utility tests, high-value UnbatchedTensor transformation tests, all TestTensorDictsBase fixture factories, a broad TestTensorDicts common-op sweep, and a tensorclass common-op sweep.
  • Keep UnbatchedTensor TensorDict-facing batch_size metadata 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 UnbatchedTensor payload 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 TestTensorDictsBase fixture factories plus common creation/transformation/modification paths for TensorDicts and tensorclasses. Under that broader coverage, the only real implementation bugs found in this pass were UnbatchedTensor TensorDict-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:

  • UnbatchedTensor values inserted into a TensorDict kept empty/stale metadata instead of adopting the parent batch size.
  • Lazy-stack retrieval of UnbatchedTensor leaves returned stale metadata, especially when source leaves had empty metadata.
  • Shape/index operations (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-facing batch_size metadata.

No public API is added; the checker is private for test/debug rollout.

  • I have raised an issue to propose this change (required for new features and bug fixes) — no issue filed for this scoped follow-up.

Runtime impact

Local measurements on this branch:

  • Representative checker benchmark over 18 transformed TensorDicts: median 6.0 us per checker call; all representative added checker calls together: median 0.11 ms per pass.
  • TestTensorDictsBase fixture checks measured locally at roughly 6.7-26.5 us per fixture construction across representative fixture types.
  • Full test/tensordict shard before adding centralized fixture checks: 7790 passed, 898 skipped in 13.69s.
  • Full test/tensordict shard after adding centralized fixture checks: 7790 passed, 898 skipped in 14.35s (real 14.94s).
  • Full test/tensordict shard 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.
  • Focused checker tests now run as 8 passed, 24030 deselected in 1.12s; the added tensorclass common-op coverage is included in test/tensorclass/test_tensorclass.py (148 passed, 1 skipped in 2.74s).

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds core functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation (update in the documentation)
  • Example (update in the folder of examples)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

Tests

  • PYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/tensordict -q
  • PYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/utils/test_utils.py -k 'check_recursive_properties' -q
  • PYTHONPATH=$PWD:$PWD/test .venv/bin/python -m pytest test/tensorclass/test_tensorclass.py -q
  • uvx 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

@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 Jun 30, 2026
@github-actions github-actions Bot added bug Something isn't working Test and removed bug Something isn't working labels Jun 30, 2026
@vmoens
vmoens force-pushed the codex/recursive-properties-checker branch from ff58b8f to 3a7c2cd Compare June 30, 2026 11:19
@github-actions github-actions Bot added CI bug Something isn't working labels Jun 30, 2026
@vmoens
vmoens force-pushed the codex/recursive-properties-checker branch from 89c98f1 to b4db4cc Compare June 30, 2026 11:42
@vmoens
vmoens merged commit 0171651 into pytorch:main Jun 30, 2026
69 of 70 checks passed
@vmoens
vmoens deleted the codex/recursive-properties-checker branch June 30, 2026 14:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working CI CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. tensorclass Test

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant