Revert D76997485 (#3183)

Summary:
Pull Request resolved: #3183

This diff reverts D76997485
https://www.internalfb.com/intern/test/562950177386883?ref_report_id=0
https://www.internalfb.com/intern/test/562950177237448?ref_report_id=0
https://www.internalfb.com/intern/test/844425154071675?ref_report_id=0
and more

Depends on D76997485

Reviewed By: PoojaAg18

Differential Revision: D78147010

fbshipit-source-id: 4b7ff779f8dbe509528b69e9020c775888818809

Author: generatedunixname89002005232357

torchrec/sparse/jagged_tensor.py

@@ -1095,18 +1095,13 @@ def _maybe_compute_stride_kjt(
     lengths: Optional[torch.Tensor],
     offsets: Optional[torch.Tensor],
     stride_per_key_per_rank: Optional[torch.IntTensor],
-    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 stride_per_key_per_rank.numel() > 0
         ):
-            # For VBE KJT, batch size should be based on inverse_indices when set.
-            if inverse_indices is not None:
-                return inverse_indices[1].shape[-1]
-
             s = stride_per_key_per_rank.sum(dim=1).max().item()
             if not torch.jit.is_scripting() and is_non_strict_exporting():
                 stride = torch.sym_int(s)
@@ -2151,7 +2146,6 @@ 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,18 +1017,6 @@ 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: