TensorDict API Additions and Indexing Semantic Fixes - #29
Merged
Conversation
added 4 commits
February 12, 2026 17:25
Handle scalar, None, and boolean indexing like PyTorch tensors and add parity tests to prevent regressions.
Consolidate and reorganize `TensorDict.__getitem__` tests into logical classes. This improves readability, maintainability, and ensures comprehensive coverage of indexing behaviors. Key changes include: - Grouping tests by functionality: `TestGetitemBasic`, `TestGetitemIsolation`, `TestGetitemZeroDim`, `TestGetitemSlicing`, `TestGetitemErrors`. - Introducing `zero_dim_td` fixture for consistent 0-dim TensorDict testing. - Parameterizing boolean and boolean tensor indexing tests for 0-dim TensorDicts. - Parameterizing slicing tests with a comprehensive list of `ALL_GETITEM_INDICES`. - Adding new tests for nested batch shape preservation and single-dimension container indexing. - Ensuring `assert_tensor_indexing_parity` includes detailed error messages.
Introduce five new TensorDict convenience methods for common container transformations: creating a shape/device-compatible instance (`like`), including keys (`select`), excluding keys (`exclude`), applying tensor functions recursively (`apply`), and top-level key remapping (`rename`). Preserve shape, device, and subclass behavior across these operations and add validation for missing keys and duplicate rename targets to keep errors explicit and predictable. Add dedicated test suites for each new method, including compiled-path checks where applicable, and expand getitem coverage with additional multidimensional/edge-case indices plus a tracked xfail for ellipsis event-dimension indexing parity.
mctigger
changed the title
added 3 commits
February 25, 2026 15:14
Drop the empty-list cat assertion that only validates PyTorch behavior and does not exercise tensor dataclass logic.
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
TensorDictconvenience APIs for structural and key-level transformations:like,select,exclude,apply, andrename.TensorContainerindexing and assignment behavior to better align withtorch.Tensorsemantics, including non-tuple indexing, ellipsis handling, and zero-dim indexing behavior.IndexType,EllipsisType) and expands test coverage for indexing and key operations.What Changed
TensorContainer indexing/setitem fixes
__getitem__/__setitem__indexing flow to consistently support torch-style index patterns.__setitem__semantics.EllipsisTypesupport for Python 3.9.TensorDict API additions
like(data, device=None)to build a new instance preserving shape/device defaults.select(*keys)to keep only selected top-level keys.exclude(*keys)to drop selected top-level keys.apply(fn)to transform all leaf tensors recursively.rename(mapping)to rename top-level keys with duplicate/missing-key validation.Tests and docs
tests/tensor_dict/test_like.pytests/tensor_dict/test_select.pytests/tensor_dict/test_exclude.pytests/tensor_dict/test_apply.pytests/tensor_dict/test_rename.pytests/tensor_dict/test_getitem.pytests/tensor_dict/test_setitem.pyTensorContainer.pyproject.tomlfor the fix release.Impact
TensorDictmanipulation APIs while preserving shape/device invariants.