Fix bug on VBE+CPU (#2256)

Summary:
Pull Request resolved: #2256

Internal users reported a bug working with VBE + CPU.   Identified regression was introduced by stray edit in D55695198.

Simple 1-line fix, but added unit test to cover this edge case for both CPU + GPU setups.

Reviewed By: TroyGarden

Differential Revision: D60430765

Author: dstaay-fb

torchrec/modules/utils.py

@@ -184,7 +184,7 @@ def _permute_tensor_by_segments(
             segment_sizes,
             tensor,
             weights,
-            tensor.numel(),
+            output_size,
         )
     return permuted_tensor, permuted_weights
 

torchrec/sparse/jagged_tensor.py

@@ -453,7 +453,7 @@ def _permute_tensor_by_segments(
             segment_sizes,
             tensor,
             weights,
-            tensor.numel(),
+            output_size,
         )
     return permuted_tensor, permuted_weights
 

torchrec/sparse/tests/test_jagged_tensor.py

@@ -1400,6 +1400,59 @@ def test_permute_vb(self) -> None:
         )
         self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
 
+    def test_permute_vb_duplicate(self) -> None:
+        values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
+        lengths = torch.IntTensor([1, 0, 1, 3, 0, 1, 0, 2, 0])
+        keys = ["index_0", "index_1", "index_2"]
+        stride_per_key_per_rank = [[2], [4], [3]]
+
+        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
+            values=values,
+            keys=keys,
+            lengths=lengths,
+            stride_per_key_per_rank=stride_per_key_per_rank,
+        )
+
+        indices = [1, 1, 0, 0, 2, 2]
+        permuted_jag_tensor = jag_tensor.permute(indices)
+
+        self.assertEqual(
+            permuted_jag_tensor.keys(),
+            ["index_1", "index_1", "index_0", "index_0", "index_2", "index_2"],
+        )
+        self.assertTrue(
+            torch.equal(
+                permuted_jag_tensor.values(),
+                torch.Tensor(
+                    [
+                        2.0,
+                        3.0,
+                        4.0,
+                        5.0,
+                        6.0,
+                        2.0,
+                        3.0,
+                        4.0,
+                        5.0,
+                        6.0,
+                        1.0,
+                        1.0,
+                        7.0,
+                        8.0,
+                        7.0,
+                        8.0,
+                    ]
+                ),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                permuted_jag_tensor.lengths(),
+                torch.IntTensor([1, 3, 0, 1, 1, 3, 0, 1, 1, 0, 1, 0, 0, 2, 0, 0, 2, 0]),
+            )
+        )
+        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
+
     def test_permute_duplicates(self) -> None:
         values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
         lengths = torch.IntTensor([0, 2, 0, 1, 1, 1, 0, 3, 0])

torchrec/sparse/tests/test_jagged_tensor_gpu.py

@@ -116,6 +116,13 @@ def test_regroup_backward(self) -> None:
         torch.allclose(actual_kt_0_grad, expected_kt_0_grad)
         torch.allclose(actual_kt_1_grad, expected_kt_1_grad)
 
+
+@skip_if_asan_class
+class TestKeyedJaggedTensorGPU(unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.device = torch.cuda.current_device()
+
     # pyre-ignore
     @unittest.skipIf(
         torch.cuda.device_count() <= 0,
@@ -187,6 +194,66 @@ def test_permute_vb(self) -> None:
         )
         self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
 
+    # pyre-ignore
+    @unittest.skipIf(
+        torch.cuda.device_count() <= 0,
+        "Not enough GPUs, this test requires at least one GPUs",
+    )
+    def test_permute_vb_duplicate(self) -> None:
+        values = torch.tensor(
+            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], device=self.device
+        )
+        lengths = torch.tensor([1, 0, 1, 3, 0, 1, 0, 2, 0], device=self.device)
+        keys = ["index_0", "index_1", "index_2"]
+        stride_per_key_per_rank = [[2], [4], [3]]
+
+        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
+            values=values,
+            keys=keys,
+            lengths=lengths,
+            stride_per_key_per_rank=stride_per_key_per_rank,
+        )
+
+        indices = [1, 1, 0, 0, 2, 2]
+        permuted_jag_tensor = jag_tensor.permute(indices)
+
+        self.assertEqual(
+            permuted_jag_tensor.keys(),
+            ["index_1", "index_1", "index_0", "index_0", "index_2", "index_2"],
+        )
+        self.assertTrue(
+            torch.equal(
+                permuted_jag_tensor.values().cpu(),
+                torch.Tensor(
+                    [
+                        2.0,
+                        3.0,
+                        4.0,
+                        5.0,
+                        6.0,
+                        2.0,
+                        3.0,
+                        4.0,
+                        5.0,
+                        6.0,
+                        1.0,
+                        1.0,
+                        7.0,
+                        8.0,
+                        7.0,
+                        8.0,
+                    ]
+                ),
+            )
+        )
+        self.assertTrue(
+            torch.equal(
+                permuted_jag_tensor.lengths().cpu(),
+                torch.IntTensor([1, 3, 0, 1, 1, 3, 0, 1, 1, 0, 1, 0, 0, 2, 0, 0, 2, 0]),
+            )
+        )
+        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
+
     # pyre-ignore
     @unittest.skipIf(
         torch.cuda.device_count() <= 0,