lazy calc offset_per_key, length_per_key (#1395)

Summary: Pull Request resolved: #1395

Differential Revision:
D49338062

Privacy Context Container: L1138451

torchrec/sparse/jagged_tensor.py

@@ -1387,6 +1387,9 @@ def _key_indices(self) -> Dict[str, int]:
         return _index_per_key
 
     def length_per_key(self) -> List[int]:
+        if self._length_per_key is not None:
+            return torch.jit._unwrap_optional(self._length_per_key)
+
         _length_per_key = _maybe_compute_length_per_key(
             keys=self._keys,
             stride=self.stride(),
@@ -1403,6 +1406,9 @@ def length_per_key_or_none(self) -> Optional[List[int]]:
         return self._length_per_key
 
     def offset_per_key(self) -> List[int]:
+        if self._offset_per_key is not None:
+            return torch.jit._unwrap_optional(self._offset_per_key)
+
         _length_per_key, _offset_per_key = _maybe_compute_offset_per_key(
             keys=self._keys,
             stride=self.stride(),