Add an interface in mc modules class to return a mask tensor. (meta-pytorch#3496)

Summary: Pull Request resolved: meta-pytorch#3496

Reviewed By: zlzhao1104

Differential Revision: D85837187

fbshipit-source-id: 34451984a6f7ec262e2aad60e6386da4fcb52514

Author: Yixin Bao

torchrec/modules/mc_embedding_modules.py

@@ -93,6 +93,19 @@ def forward(
             return embedding_res, None
         return embedding_res, features
 
+    def lookup_remapped_lengths_mask(
+        self,
+        features: KeyedJaggedTensor,
+    ) -> torch.Tensor:
+        features = self._managed_collision_collection(features)
+        remapped_lengths = return_remapped_lengths_as_mask(features)
+        return remapped_lengths
+
+
+@torch.fx.wrap
+def return_remapped_lengths_as_mask(features: KeyedJaggedTensor) -> torch.Tensor:
+    return features.lengths().to(torch.bool)
+
 
 class ManagedCollisionEmbeddingCollection(BaseManagedCollisionEmbeddingCollection):
     """

torchrec/modules/tests/test_mc_embedding_modules.py

@@ -21,6 +21,7 @@
 from torchrec.modules.mc_embedding_modules import (
     ManagedCollisionEmbeddingBagCollection,
     ManagedCollisionEmbeddingCollection,
+    return_remapped_lengths_as_mask,
 )
 from torchrec.modules.mc_modules import (
     DistanceLFU_EvictionPolicy,
@@ -409,3 +410,20 @@ def test_mc_collection_traceable(self) -> None:
         )
         mcc.train(False)
         symbolic_trace(mcc, leaf_modules=[ComputeJTDictToKJT.__name__])
+
+    def test_return_remapped_lengths_as_mask(self) -> None:
+        mask = return_remapped_lengths_as_mask(
+            KeyedJaggedTensor(
+                keys=["f0"],
+                values=torch.rand(6),
+                lengths=torch.tensor([1, 0, 1, 0, 1, 0, 0, 1, 1, 1], dtype=torch.int64),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                mask,
+                torch.tensor(
+                    [True, False, True, False, True, False, False, True, True, True]
+                ),
+            )
+        )