Add max_length to JaggedTensor api (#3160)

Summary:
Pull Request resolved: #3160

We often need to recompute the max_length of a JaggedTensor on the fly, which makes the compute graph fragmented. This change follows the caching approach of existing lengths/offsets implementation for max_length.

Reviewed By: shruthign

Differential Revision: D77709565

fbshipit-source-id: cad35e4c16607c859291d18d1b4f90f98d8a0002

Author: yunjiangster

torchrec/sparse/jagged_tensor.py

@@ -98,6 +98,15 @@ def _maybe_compute_lengths(
     return lengths
 
 
+@torch.fx.wrap
+def _maybe_compute_max_length(lengths: torch.Tensor, max_length: Optional[int]) -> int:
+    if max_length is None:
+        if lengths.numel() == 0:
+            return 0
+        max_length = int(lengths.max().item())
+    return max_length
+
+
 def _maybe_compute_offsets(
     lengths: Optional[torch.Tensor], offsets: Optional[torch.Tensor]
 ) -> torch.Tensor:
@@ -581,14 +590,15 @@ class JaggedTensor(Pipelineable, metaclass=JaggedTensorMeta):
             offsets.
     """
 
-    _fields = ["_values", "_weights", "_lengths", "_offsets"]
+    _fields = ["_values", "_weights", "_lengths", "_offsets", "_max_length"]
 
     def __init__(
         self,
         values: torch.Tensor,
         weights: Optional[torch.Tensor] = None,
         lengths: Optional[torch.Tensor] = None,
         offsets: Optional[torch.Tensor] = None,
+        max_length: Optional[int] = None,
     ) -> None:
 
         self._values: torch.Tensor = values
@@ -600,6 +610,7 @@ def __init__(
             _assert_tensor_has_no_elements_or_has_integers(lengths, "lengths")
         self._lengths: Optional[torch.Tensor] = lengths
         self._offsets: Optional[torch.Tensor] = offsets
+        self._max_length: Optional[int] = max_length
 
     @staticmethod
     def empty(
@@ -630,6 +641,7 @@ def empty(
             offsets=torch.empty(0, dtype=lengths_dtype, device=device),
             lengths=torch.empty(0, dtype=lengths_dtype, device=device),
             weights=weights,
+            max_length=0,
         )
 
     @staticmethod
@@ -912,6 +924,26 @@ def lengths_or_none(self) -> Optional[torch.Tensor]:
         """
         return self._lengths
 
+    def max_length(self) -> int:
+        """
+        Get the maximum length of the JaggedTensor.
+
+        Returns:
+            int: the maximum length of the JaggedTensor.
+        """
+        _max_length = _maybe_compute_max_length(self.lengths(), self._max_length)
+        self._max_length = _max_length
+        return _max_length
+
+    def max_length_or_none(self) -> Optional[int]:
+        """
+        Get the maximum length of the JaggedTensor. If not computed, return None.
+
+        Returns:
+            Optional[int]: the maximum length of the JaggedTensor.
+        """
+        return self._max_length
+
     def offsets(self) -> torch.Tensor:
         """
         Get JaggedTensor offsets. If not computed, compute it from lengths.
@@ -973,6 +1005,7 @@ def to(self, device: torch.device, non_blocking: bool = False) -> "JaggedTensor"
         weights = self._weights
         lengths = self._lengths
         offsets = self._offsets
+        max_length = self._max_length
         return JaggedTensor(
             values=self._values.to(device, non_blocking=non_blocking),
             weights=(
@@ -990,6 +1023,7 @@ def to(self, device: torch.device, non_blocking: bool = False) -> "JaggedTensor"
                 if offsets is not None
                 else None
             ),
+            max_length=max_length,
         )
 
     @torch.jit.unused

torchrec/sparse/tests/test_jagged_tensor.py

@@ -568,6 +568,14 @@ def test_length_vs_offset(self) -> None:
         self.assertTrue(torch.equal(j_offset.lengths(), j_lens.lengths()))
         self.assertTrue(torch.equal(j_offset.offsets(), j_lens.offsets().int()))
 
+    def test_max_length(self) -> None:
+        values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
+        offsets = torch.IntTensor([0, 2, 2, 3, 4, 5, 8])
+        jt = JaggedTensor(values=values, offsets=offsets)
+        self.assertIsNone(jt.max_length_or_none())
+        self.assertEqual(jt.max_length(), 3)
+        self.assertEqual(jt.max_length_or_none(), 3)
+
     def test_empty(self) -> None:
         jt = JaggedTensor.empty(values_dtype=torch.int64)