From 9139195b91577965c3a4be9c0fe66e8b3ff98c5a Mon Sep 17 00:00:00 2001 From: Akshay Shekher Date: Thu, 23 Oct 2025 10:29:22 -0700 Subject: [PATCH] Sync tensordict (#3478) Summary: Pull Request resolved: https://github.com/meta-pytorch/torchrec/pull/3478 Imported commits from public tensordict repo using ``` python3 pytorch/import.py --project_name tensordict --no_submit ``` Manual Change: * adding a new dependency: pyvers * imported a new dependency pyvers in the stack * some import failures due to checked in code missing new lines * hypothesis: code formatter removed the new lines * manually added the new lines back and re-ran the script * added pyi files to library source to fix linter errors * disabled tests that use ray because we can't initialize a ray cluster on sandcastle hosts Differential Revision: D85157386 --- torchrec/distributed/embedding.py | 2 +- torchrec/distributed/embeddingbag.py | 2 +- torchrec/distributed/test_utils/test_model.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/torchrec/distributed/embedding.py b/torchrec/distributed/embedding.py index d8554edea..9ddb6e91c 100644 --- a/torchrec/distributed/embedding.py +++ b/torchrec/distributed/embedding.py @@ -1469,7 +1469,7 @@ def input_dist( ) -> Awaitable[Awaitable[KJTList]]: need_permute: bool = True if isinstance(features, TensorDict): - feature_keys = list(features.keys()) # pyre-ignore[6] + feature_keys = list(features.keys()) if self._features_order: feature_keys = [feature_keys[i] for i in self._features_order] need_permute = False diff --git a/torchrec/distributed/embeddingbag.py b/torchrec/distributed/embeddingbag.py index fd6117884..4375baa4f 100644 --- a/torchrec/distributed/embeddingbag.py +++ b/torchrec/distributed/embeddingbag.py @@ -1513,7 +1513,7 @@ def input_dist( in advance """ if isinstance(features, TensorDict): - feature_keys = list(features.keys()) # pyre-ignore[6] + feature_keys = list(features.keys()) if len(self._features_order) > 0: feature_keys = [feature_keys[i] for i in self._features_order] self._has_features_permute = False # feature_keys are in order diff --git a/torchrec/distributed/test_utils/test_model.py b/torchrec/distributed/test_utils/test_model.py index cb7004670..d0d0c7a8d 100644 --- a/torchrec/distributed/test_utils/test_model.py +++ b/torchrec/distributed/test_utils/test_model.py @@ -298,7 +298,7 @@ def _validate_pooling_factor( idlist_features, global_idlist_indices, global_idlist_lengths ) } - global_idlist_input = TensorDict(source=dict_of_nt) + global_idlist_input = TensorDict(source=dict_of_nt) # pyre-ignore[6] assert ( len(idscore_features) == 0 @@ -409,7 +409,7 @@ def _validate_pooling_factor( local_idlist_lengths, ) } - local_idlist_input = TensorDict(source=dict_of_nt) + local_idlist_input = TensorDict(source=dict_of_nt) # pyre-ignore[6] assert ( len(idscore_features) == 0 ), "TensorDict does not support weighted features"