Add option to pass inverse indices to new KJT after permute

Summary: allow inverse indices to be included in new KJT after permute. For modifying vbe KJT outside of input dist.

Differential Revision: D53732141

torchrec/sparse/jagged_tensor.py

@@ -1645,7 +1645,10 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
         return split_list
 
     def permute(
-        self, indices: List[int], indices_tensor: Optional[torch.Tensor] = None
+        self,
+        indices: List[int],
+        indices_tensor: Optional[torch.Tensor] = None,
+        include_inverse_indices: bool = False,
     ) -> "KeyedJaggedTensor":
 
         if indices_tensor is None:
@@ -1715,7 +1718,9 @@ def permute(
             offset_per_key=None,
             index_per_key=None,
             jt_dict=None,
-            inverse_indices=None,
+            inverse_indices=self.inverse_indices_or_none()
+            if include_inverse_indices
+            else None,
         )
         return kjt