jagged_tensor minor refactoring (#2860)

Summary:
Pull Request resolved: #2860

# context
* split test_jagged_tensor into multiple files to increase readability
* use hypothesis.strategies to generate the test cases (the standard approach)
* add some comments in jagged_tensor.py
* use `torch.diff` to replace `_to_lengths` function.

Reviewed By: ezyang, aporialiao

Differential Revision: D56177133

fbshipit-source-id: 245536e7a1d10eee132a4e87b5813840249b793b

Author: TroyGarden

torchrec/sparse/jagged_tensor.py

@@ -52,6 +52,9 @@
 
 
 def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
+    """
+    moving a tensor from cpu to cuda using pinned memory (non_blocking) is generally faster
+    """
     if is_torchdynamo_compiling():
         # TODO: remove once FakeTensor supports pin_memory() and to(..., non_blocking=True)
         return tensor.to(device=device)
@@ -64,6 +67,9 @@ def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
 
 
 def _cumsum(o: List[int]) -> List[int]:
+    """
+    python-list version of converting lengths --> offsets
+    """
     ret = [0] * (len(o) + 1)
     for i in range(len(o)):
         ret[i + 1] = ret[i] + o[i]
@@ -92,7 +98,7 @@ def _maybe_compute_lengths(
 ) -> torch.Tensor:
     if lengths is None:
         assert offsets is not None
-        lengths = _to_lengths(offsets)
+        lengths = torch.diff(offsets)
     return lengths