Implement narrow from a regular tensor to jagged tensor

torch/nested/__init__.py

@@ -1,7 +1,7 @@
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import torch
-from torch import Tensor
+from torch import SymInt, Tensor
 from torch._C import _add_docstr, _nested  # type: ignore[attr-defined]
 
 from torch.types import _device as Device, _dtype as DType
@@ -186,3 +186,59 @@ def nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires
         return nt
     else:
         raise RuntimeError(f"Specified layout is unsupported for nested tensors: {layout}")
+
+
+def narrow(tensor: Tensor, dim: int, start: Union[int, Tensor], length: Union[int, Tensor], layout=torch.jagged):
+    r"""
+Constructs a nested tensor (which might be a view) from :attr:`tensor`, a strided tensor. This follows
+similar semantics ot torch.Tensor.narrow, where in the :attr:`dim`-th dimension the new nested tensor
+(maybe view) shows only the elements in the interval `[start, start+length]`. As nested representations
+allow for a different `start` and `length` at each 'row' of that dimension, :attr:`start` and :attr:`length`
+can also be tensors of shape `tensor.shape[0] x 1`.
+
+
+Args:
+    tensor (:class:`torch.Tensor`): a strided tensor, which will be used as the underlying data
+    for the nested tensor if using the jagged layout or will be copied for the strided layout.
+    dim (int): the dimension where narrow will be applied. Only `dim=1` is supported for the
+    jagged layout, while strided supports all dim
+    start (Union[int, :class:`torch.Tensor`]): starting element for the narrow operation
+    length (Union[int, :class:`torch.Tensor`]): number of elements taken during the narrow op
+
+Keyword arguments:
+    layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
+        Only strided and jagged layouts are supported. Default: if None, the jagged layout.
+
+Example::
+
+    >>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
+    >>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
+    >>> nt = torch.nested.nested_tensor([a, b], requires_grad=True)
+    >>> nt.is_leaf
+    True
+    """
+    if not isinstance(start, int) or isinstance(start, Tensor):
+        raise RuntimeError("start must be an integer or a tensor")
+
+    if not isinstance(length, int) or isinstance(length, Tensor):
+        raise RuntimeError("length must be an integer or a tensor")
+
+    if layout == torch.strided:
+        nt = as_nested_tensor(torch.unbind(tensor), layout=torch.strided).narrow(dim, start, length)
+    elif layout == torch.jagged:
+        if dim != 1:
+            raise RuntimeError("jagged layout only supports dim=1")
+
+        from torch.nested._internal.nested_tensor import jagged_from_tensor_and_lengths
+
+        if isinstance(start, int):
+            start = torch.tensor([start], device=tensor.device, dtype=torch.int64)
+
+        if isinstance(length, int):
+            length = torch.tensor([length], device=tensor.device, dtype=torch.int64)
+
+        nt, _, _ = jagged_from_tensor_and_lengths(tensor, start, length)
+    else:
+        raise RuntimeError(f"Specified layout is unsupported for nested narrow: {layout}")
+
+    return nt

torch/nested/_internal/nested_tensor.py

@@ -2,6 +2,7 @@
 
 import torch
 from torch._C import DispatchKey, DispatchKeySet
+from torch._prims_common import is_expandable_to
 from torch.fx.experimental.symbolic_shapes import free_symbols
 from torch.utils.weak import WeakTensorKeyDictionary
 from typing import *  # noqa: F403
@@ -21,6 +22,7 @@ def get_tensor_id(tensor, *, coeff=1):
 class NestedTensor(torch.Tensor):
     _values: torch.Tensor  # type: ignore[assignment]
     _offsets: torch.Tensor
+    _lengths: Optional[torch.Tensor]
     # NOTE [ Singleton ints for ragged sizes and strides ]
     #
     # Jagged layout tensors are tensors that represent a n-dim tensor with a
@@ -46,6 +48,7 @@ def __new__(
         values,
         offsets,
         *,
+        lengths=None,
         ragged_size=None,
         **kwargs,
     ):
@@ -69,7 +72,7 @@ def __new__(
         )
         return r
 
-    def __init__(self, values, offsets, *, ragged_size=None, **kwargs):
+    def __init__(self, values, offsets, *, lengths=None, ragged_size=None, **kwargs):
         super().__init__()
         # Only support jagged for now.
         assert offsets is not None
@@ -97,21 +100,25 @@ def __init__(self, values, offsets, *, ragged_size=None, **kwargs):
             )
         self._values = values
         self._offsets = offsets
+        self._lengths = lengths
 
     def values(self):
         return self._values
 
     def offsets(self):
         return self._offsets
 
+    def lengths(self):
+        return self._lengths
+
     def __repr__(self):
         # We should implement this in torch/_tensor_str.py instead
         grad_fn_str = (
             f", requires_grad={self.requires_grad}" if self.requires_grad else ""
         )
         if self.grad_fn:
             grad_fn_str = f", grad_fn={self.grad_fn}"
-        return f"NestedTensor(size={self._size}, offsets={self._offsets}{grad_fn_str})"
+        return f"NestedTensor(size={self._size}, offsets={self._offsets}{grad_fn_str}, contiguous={self._lengths is None})"
 
     def __reduce_ex__(self, proto):
         state = torch._utils._get_obj_state(self)
@@ -131,13 +138,14 @@ def __tensor_flatten__(self):
             "requires_grad": self.requires_grad,
             "ragged_size": self._size[self._ragged_idx],
         }
-        return ["_values", "_offsets"], ctx
+        return ["_values", "_offsets", "_lengths"], ctx
 
     @staticmethod
     def __tensor_unflatten__(inner_tensors: Dict, meta):
         assert len(inner_tensors) == 2
         values = inner_tensors["_values"]
         offsets = inner_tensors["_offsets"]
+        lengths = inner_tensors["_lengths"]
 
         # NOTE [ Storing symbolic values as plain attributes on subclasses ]
         #
@@ -173,6 +181,7 @@ def __tensor_unflatten__(inner_tensors: Dict, meta):
         return NestedTensor(
             values,
             offsets=offsets,
+            lengths=lengths,
             ragged_size=meta["ragged_size"],
             requires_grad=meta["requires_grad"],
         )
@@ -232,6 +241,17 @@ def backward(ctx, gO: NestedTensor):  # type: ignore[override]
         return gO.values(), None, None
 
 
+# Not actually a view!
+class ViewNonContiguousNestedFromBuffer(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values: torch.Tensor, offsets: torch.Tensor, lengths: torch.Tensor):  # type: ignore[override]
+        return NestedTensor(values.detach(), offsets=offsets, lengths=lengths)
+
+    @staticmethod
+    def backward(ctx, gO: NestedTensor):  # type: ignore[override]
+        return gO.values(), None, None
+
+
 # Need to make it obvious that users should be passing in offsets
 def jagged_from_list(
     tensors: List[torch.Tensor],
@@ -285,5 +305,41 @@ def jagged_from_list(
     return ViewNestedFromBuffer.apply(values, offsets), offsets  # type: ignore[call-overload]
 
 
+def jagged_from_tensor_and_lengths(
+    tensor: torch.Tensor, starts: torch.Tensor, lengths: torch.Tensor
+) -> Tuple[NestedTensor, torch.Tensor, torch.Tensor]:
+    """Constructs a NestedTensor backed by jagged layout from a tensor, starts of sequences, and sequence lengths"""
+    batch_size = tensor.shape[0]
+    if is_expandable_to(starts.shape, (batch_size, 1)) and is_expandable_to(
+        lengths.shape, (batch_size, 1)
+    ):
+        start_list = starts.expand(batch_size, 1)
+        length_list = lengths.expand(batch_size, 1)
+    else:
+        raise RuntimeError(
+            "When constructing a jagged nested tensor using narrow(), "
+            "your start and length must be a Tensor that broadcasts to input.shape[0] x 1"
+        )
+
+    # Calculate jagged offsets
+    assert (
+        len(tensor.shape) >= 2
+    ), "tensor must at least be 2D for the nested narrow op to work"
+    max_seq_len = tensor.shape[1]
+    offset_lengths = max_seq_len * torch.arange(
+        0, batch_size, dtype=torch.int64, device=tensor.device
+    )
+    # Jagged layout specifies that offsets are stored as int64 on the same device as values.
+    offsets = start_list + offset_lengths
+
+    # Reshape buffer to flatten the 1st and 2nd dimension
+    if len(tensor.shape) > 2:
+        values = tensor.reshape(-1, *tensor.shape[2:])
+    else:
+        values = tensor.reshape(-1)
+
+    return ViewNonContiguousNestedFromBuffer.apply(values, offsets, length_list), offsets, length_list  # type: ignore[call-overload]
+
+
 def buffer_from_jagged(jagged):
     return ViewBufferFromNested.apply(jagged)

torch/nested/_internal/ops.py

@@ -36,7 +36,8 @@ def check_schema(schema_str: str, func, *args, **kwargs) -> None:
 
     arg_type_check_fns = {
         "t": lambda x: isinstance(x, torch.Tensor) and not isinstance(x, NestedTensor),
-        "jt": lambda x: isinstance(x, NestedTensor),
+        "jt": lambda x: isinstance(x, NestedTensor) and x._lengths is None,
+        "jt_nc": lambda x: isinstance(x, NestedTensor),
         "any": lambda x: True,
     }
     for i, named_arg_type in enumerate(named_arg_types):