Update KJT stride calculation logic to be based off of inverse_indices for VBE KJTs. (#2949)

Summary:

Update the `_maybe_compute_stride_kjt` logic to calculate stride based off of `inverse_indices` for VBE KJTs.

Currently, stride of VBE KJT with `stride_per_key_per_rank` is calculated as the max "stride per key". This is different from the batch size of the EBC output KeyedTensor which is based off of inverse_indices. This causes issues in IR module serialization: debug doc.

Differential Revision: D74273083

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1097,11 +1097,15 @@ def _maybe_compute_stride_kjt(
     lengths: Optional[torch.Tensor],
     offsets: Optional[torch.Tensor],
     stride_per_key_per_rank: Optional[List[List[int]]],
+    inverse_indices: Optional[Tuple[List[str], torch.Tensor]] = None,
 ) -> int:
     if stride is None:
         if len(keys) == 0:
             stride = 0
         elif stride_per_key_per_rank is not None and len(stride_per_key_per_rank) > 0:
+            # For VBE KJT, use inverse_indices for the batch size of the EBC output KeyedTensor.
+            if inverse_indices is not None and inverse_indices[1].numel() > 0:
+                return inverse_indices[1].shape[-1]
             stride = max([sum(s) for s in stride_per_key_per_rank])
         elif offsets is not None and offsets.numel() > 0:
             stride = (offsets.numel() - 1) // len(keys)
@@ -2165,6 +2169,7 @@ def stride(self) -> int:
             self._lengths,
             self._offsets,
             self._stride_per_key_per_rank,
+            self._inverse_indices,
         )
         self._stride = stride
         return stride

torchrec/sparse/tests/test_keyed_jagged_tensor.py

@@ -1017,6 +1017,18 @@ def test_meta_device_compatibility(self) -> None:
             lengths=torch.tensor([], device=torch.device("meta")),
         )
 
+    def test_vbe_kjt_stride(self) -> None:
+        inverse_indices = torch.tensor([[0, 1, 0], [0, 0, 0]])
+        kjt = KeyedJaggedTensor(
+            keys=["f1", "f2", "f3"],
+            values=torch.tensor([5, 6, 7, 1, 2, 3, 0, 1]),
+            lengths=torch.tensor([3, 3, 2]),
+            stride_per_key_per_rank=[[2], [1]],
+            inverse_indices=(["f1", "f2"], inverse_indices),
+        )
+
+        self.assertEqual(kjt.stride(), inverse_indices.shape[1])
+
 
 class TestKeyedJaggedTensorScripting(unittest.TestCase):
     def test_scriptable_forward(self) -> None: