Create a new tensor arg for stride_per_key_per_rank to facilitate torch.export

Summary:
# Context
* Currently torchrec IR serializer can't handle variable batch KJT use case.
* To support VBE KJT, the `stride_per_key_per_rank` field needs to be flattened as a variable in the pytree flatten spec for a VBE KJT to be unflattened correctly by`torch.export.
* Currently `stride_per_key_per_rank` is a List. To flatten the `stride_per_key_per_rank`  info as a variable we have to add a new tensor field for it.

# Ref

Differential Revision: D74207283

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1778,6 +1778,7 @@ def __init__(
         index_per_key: Optional[Dict[str, int]] = None,
         jt_dict: Optional[Dict[str, JaggedTensor]] = None,
         inverse_indices: Optional[Tuple[List[str], torch.Tensor]] = None,
+        stride_per_key_per_rank_tensor: Optional[torch.Tensor] = None,
     ) -> None:
         """
         This is the constructor for KeyedJaggedTensor is jit.scriptable and PT2 compatible.
@@ -1794,6 +1795,11 @@ def __init__(
         self._stride_per_key_per_rank: Optional[List[List[int]]] = (
             stride_per_key_per_rank
         )
+
+        self._stride_per_key_per_rank_tensor: torch.Tensor = torch.empty(0)
+        if stride_per_key_per_rank_tensor is not None:
+            self._stride_per_key_per_rank_tensor = stride_per_key_per_rank_tensor
+
         self._stride_per_key: Optional[List[int]] = stride_per_key
         self._length_per_key: Optional[List[int]] = length_per_key
         self._offset_per_key: Optional[List[int]] = offset_per_key
@@ -2183,7 +2189,7 @@ def stride_per_key(self) -> List[int]:
         """
         stride_per_key = _maybe_compute_stride_per_key(
             self._stride_per_key,
-            self._stride_per_key_per_rank,
+            self._stride_per_key_per_rank_optional,
             self.stride(),
             self._keys,
         )
@@ -2198,7 +2204,27 @@ def stride_per_key_per_rank(self) -> List[List[int]]:
             List[List[int]]: stride per key per rank of the KeyedJaggedTensor.
         """
         stride_per_key_per_rank = self._stride_per_key_per_rank
-        return stride_per_key_per_rank if stride_per_key_per_rank is not None else []
+
+        if stride_per_key_per_rank is not None:
+            return stride_per_key_per_rank
+
+        if self._stride_per_key_per_rank_tensor.numel() > 0:
+            return self._stride_per_key_per_rank_tensor.tolist()
+
+        return []
+
+    @property
+    def _stride_per_key_per_rank_optional(self) -> Optional[List[List[int]]]:
+        if self._stride_per_key_per_rank is not None:
+            return self._stride_per_key_per_rank
+
+        if self._stride_per_key_per_rank_tensor.numel() > 0:
+            stride_per_key_per_rank: List[List[int]] = (
+                self._stride_per_key_per_rank_tensor.tolist()
+            )
+            return stride_per_key_per_rank
+
+        return None
 
     def variable_stride_per_key(self) -> bool:
         """
@@ -2209,7 +2235,7 @@ def variable_stride_per_key(self) -> bool:
         """
         if self._variable_stride_per_key is not None:
             return self._variable_stride_per_key
-        return self._stride_per_key_per_rank is not None
+        return self._stride_per_key_per_rank_optional is not None
 
     def inverse_indices(self) -> Tuple[List[str], torch.Tensor]:
         """
@@ -2374,6 +2400,7 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
                         index_per_key=self._index_per_key,
                         jt_dict=self._jt_dict,
                         inverse_indices=None,
+                        stride_per_key_per_rank_tensor=None,
                     )
                 )
             elif segment == 0:
@@ -2410,6 +2437,7 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
                         index_per_key=None,
                         jt_dict=None,
                         inverse_indices=None,
+                        stride_per_key_per_rank_tensor=None,
                     )
                 )
             else:
@@ -2456,6 +2484,7 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
                             index_per_key=None,
                             jt_dict=None,
                             inverse_indices=None,
+                            stride_per_key_per_rank_tensor=None,
                         )
                     )
                 else:
@@ -2492,6 +2521,7 @@ def split(self, segments: List[int]) -> List["KeyedJaggedTensor"]:
                             index_per_key=None,
                             jt_dict=None,
                             inverse_indices=None,
+                            stride_per_key_per_rank_tensor=None,
                         )
                     )
             start = end
@@ -2598,12 +2628,15 @@ def permute(
             index_per_key=None,
             jt_dict=None,
             inverse_indices=None,
+            stride_per_key_per_rank_tensor=None,
         )
         return kjt
 
     def flatten_lengths(self) -> "KeyedJaggedTensor":
         stride_per_key_per_rank = (
-            self._stride_per_key_per_rank if self.variable_stride_per_key() else None
+            self._stride_per_key_per_rank_optional
+            if self.variable_stride_per_key()
+            else None
         )
         return KeyedJaggedTensor(
             keys=self._keys,
@@ -2620,6 +2653,7 @@ def flatten_lengths(self) -> "KeyedJaggedTensor":
             index_per_key=None,
             jt_dict=None,
             inverse_indices=None,
+            stride_per_key_per_rank_tensor=None,
         )
 
     def __getitem__(self, key: str) -> JaggedTensor:
@@ -2759,7 +2793,9 @@ def to(
         lengths = self._lengths
         offsets = self._offsets
         stride_per_key_per_rank = (
-            self._stride_per_key_per_rank if self.variable_stride_per_key() else None
+            self._stride_per_key_per_rank_optional
+            if self.variable_stride_per_key()
+            else None
         )
         length_per_key = self._length_per_key
         lengths_offset_per_key = self._lengths_offset_per_key
@@ -2804,6 +2840,7 @@ def to(
             index_per_key=index_per_key,
             jt_dict=jt_dict,
             inverse_indices=inverse_indices,
+            stride_per_key_per_rank_tensor=None,
         )
 
     def __str__(self) -> str:
@@ -2835,7 +2872,9 @@ def pin_memory(self) -> "KeyedJaggedTensor":
         lengths = self._lengths
         offsets = self._offsets
         stride_per_key_per_rank = (
-            self._stride_per_key_per_rank if self.variable_stride_per_key() else None
+            self._stride_per_key_per_rank_optional
+            if self.variable_stride_per_key()
+            else None
         )
         inverse_indices = self._inverse_indices
         if inverse_indices is not None:
@@ -2856,6 +2895,7 @@ def pin_memory(self) -> "KeyedJaggedTensor":
             index_per_key=self._index_per_key,
             jt_dict=None,
             inverse_indices=inverse_indices,
+            stride_per_key_per_rank_tensor=None,
         )
 
     def dist_labels(self) -> List[str]:

torchrec/sparse/tests/test_keyed_jagged_tensor.py

@@ -1018,16 +1018,32 @@ def test_meta_device_compatibility(self) -> None:
         )
 
     def test_vbe_kjt_stride(self) -> None:
+        stride_per_key_per_rank = [[2], [1]]
         inverse_indices = torch.tensor([[0, 1, 0], [0, 0, 0]])
-        kjt = KeyedJaggedTensor(
+        kjt_1 = 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]],
+            stride_per_key_per_rank=stride_per_key_per_rank,
+            inverse_indices=(["f1", "f2"], inverse_indices),
+        )
+        kjt_2 = 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_tensor=torch.tensor(stride_per_key_per_rank),
             inverse_indices=(["f1", "f2"], inverse_indices),
         )
 
-        self.assertEqual(kjt.stride(), inverse_indices.shape[1])
+        self.assertEqual(kjt_1.stride(), inverse_indices.shape[1])
+        self.assertEqual(kjt_1.stride_per_key_per_rank(), stride_per_key_per_rank)
+        self.assertEqual(
+            kjt_1._stride_per_key_per_rank_optional, stride_per_key_per_rank
+        )
+        self.assertEqual(kjt_2.stride_per_key_per_rank(), stride_per_key_per_rank)
+        self.assertEqual(
+            kjt_2._stride_per_key_per_rank_optional, stride_per_key_per_rank
+        )
 
 
 class TestKeyedJaggedTensorScripting(unittest.TestCase):