Back out "JaggedTensor permute - less CPU ops" (#2862)

Summary:
Pull Request resolved: #2862

Original commit changeset: 257a9a45b204

Original Phabricator Diff: D70609204

The original diff causes our integration tests failing: https://fburl.com/mlhub/q87huj8u

Example of failed test job: f714668378

The error message is "AssertionError: inverse indices must be provided from KJT if using variable batch size per feature."

Backout this diff to unblock our tests

Reviewed By: dshi7

Differential Revision: D72187365

fbshipit-source-id: 58ff117c3583af7427eee2193cb510e3e0ea79dc

Author: Microve

torchrec/sparse/jagged_tensor.py

@@ -2517,23 +2517,22 @@ def permute(
         permuted_stride_per_key_per_rank: List[List[int]] = []
         permuted_length_per_key: List[int] = []
         permuted_length_per_key_sum = 0
-        keys = self._keys
-        variable_stride_per_key = self.variable_stride_per_key()
-        stride_per_key_per_rank = self.stride_per_key_per_rank()
         for index in indices:
-            key = keys[index]
+            key = self.keys()[index]
             permuted_keys.append(key)
             permuted_length_per_key.append(length_per_key[index])
-            if variable_stride_per_key:
-                permuted_stride_per_key_per_rank.append(stride_per_key_per_rank[index])
+            if self.variable_stride_per_key():
+                permuted_stride_per_key_per_rank.append(
+                    self.stride_per_key_per_rank()[index]
+                )
 
         permuted_length_per_key_sum = sum(permuted_length_per_key)
         if not torch.jit.is_scripting() and is_non_strict_exporting():
             torch._check_is_size(permuted_length_per_key_sum)
             torch._check(permuted_length_per_key_sum != -1)
             torch._check(permuted_length_per_key_sum != 0)
 
-        if variable_stride_per_key:
+        if self.variable_stride_per_key():
             length_per_key_tensor = _pin_and_move(
                 torch.tensor(self.length_per_key()), self.device()
             )
@@ -2578,7 +2577,7 @@ def permute(
                 permuted_length_per_key_sum,
             )
         stride_per_key_per_rank = (
-            permuted_stride_per_key_per_rank if variable_stride_per_key else None
+            permuted_stride_per_key_per_rank if self.variable_stride_per_key() else None
         )
         kjt = KeyedJaggedTensor(
             keys=permuted_keys,