Update the _maybe_compute_stride_kjt logic to calculate stride based off of inverse_indices for VBE KJTs. (#2925)

Summary:

See D73051959 for context.

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](https://docs.google.com/document/d/1yQhI484cgVloSqIBPAeTQhzfb3ltjvMRiLaQDceHGOU/edit?tab=t.0#heading=h.c66chahhl8df).

Differential Revision: D73824764

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1099,10 +1099,13 @@ 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]],
 ) -> int:
     if stride is None:
         if len(keys) == 0:
             stride = 0
+        elif inverse_indices is not None and inverse_indices[1].numel() > 0:
+            return inverse_indices[1].shape[1]
         elif stride_per_key_per_rank is not None and len(stride_per_key_per_rank) > 0:
             stride = max([sum(s) for s in stride_per_key_per_rank])
         elif offsets is not None and offsets.numel() > 0:
@@ -2174,6 +2177,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

@@ -1033,6 +1033,7 @@ def test_flatten_unflatten_with_vbe(self) -> None:
             kjt.stride_per_key_per_rank(), unflattened_kjt.stride_per_key_per_rank()
         )
         self.assertEqual(kjt.inverse_indices(), unflattened_kjt.inverse_indices())
+        self.assertEqual(kjt.stride(), kjt.inverse_indices()[1].shape[1])
 
 
 class TestKeyedJaggedTensorScripting(unittest.TestCase):