Test include stride_per_key to KJT's flatten and unflatten

Differential Revision: D73051959

Author: jd7-tr

torchrec/sparse/jagged_tensor.py

@@ -1756,6 +1756,7 @@ class KeyedJaggedTensor(Pipelineable, metaclass=JaggedTensorMeta):
         "_weights",
         "_lengths",
         "_offsets",
+        "_stride_per_key",
     ]
 
     def __init__(
@@ -3033,7 +3034,17 @@ def dist_init(
 def _kjt_flatten(
     t: KeyedJaggedTensor,
 ) -> Tuple[List[Optional[torch.Tensor]], List[str]]:
-    return [getattr(t, a) for a in KeyedJaggedTensor._fields], t._keys
+    values, keys = [getattr(t, a) for a in KeyedJaggedTensor._fields[:-1]], t._keys
+
+    # append stride_per_key to the tensors
+    stride_per_key = t.stride_per_key()
+    values.append(
+        torch.tensor(stride_per_key, dtype=torch.int32, device=t.device())
+        if stride_per_key is not None
+        else None
+    )
+
+    return values, keys
 
 
 def _kjt_flatten_with_keys(
@@ -3049,13 +3060,22 @@ def _kjt_flatten_with_keys(
 def _kjt_unflatten(
     values: List[Optional[torch.Tensor]], context: List[str]  # context is the _keys
 ) -> KeyedJaggedTensor:
-    return KeyedJaggedTensor(context, *values)
+    return KeyedJaggedTensor(context, *values[:-1], stride_per_key=values[-1])
 
 
 def _kjt_flatten_spec(
     t: KeyedJaggedTensor, spec: TreeSpec
 ) -> List[Optional[torch.Tensor]]:
-    return [getattr(t, a) for a in KeyedJaggedTensor._fields]
+    stride_per_key = t.stride_per_key()
+    stride_per_key_tensor: Optional[torch.Tensor] = (
+        torch.tensor(stride_per_key, dtype=torch.int32, device=t.device())
+        if stride_per_key is not None
+        else None
+    )
+    values = [getattr(t, a) for a in KeyedJaggedTensor._fields[:-1]]
+    values.append(stride_per_key_tensor)
+
+    return values
 
 
 register_pytree_node(

torchrec/sparse/tests/test_keyed_jagged_tensor.py

@@ -1017,6 +1017,22 @@ def test_meta_device_compatibility(self) -> None:
             lengths=torch.tensor([], device=torch.device("meta")),
         )
 
+    def test_flatten_unflatten_with_vbe(self) -> None:
+        kjt = KeyedJaggedTensor(
+            keys=["f1", "f2"],
+            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"], torch.tensor([[0, 1, 0], [0, 0, 0]])),
+        )
+
+        flat_kjt, spec = pytree.tree_flatten(kjt)
+        unflattened_kjt = pytree.tree_unflatten(flat_kjt, spec)
+
+        self.assertEqual(
+            kjt.stride_per_key(), unflattened_kjt.stride_per_key().tolist()
+        )
+
 
 class TestKeyedJaggedTensorScripting(unittest.TestCase):
     def test_scriptable_forward(self) -> None: