254 lines
11 KiB
Python
254 lines
11 KiB
Python
from typing import List, Optional, Union, Sequence
|
|
|
|
import torch
|
|
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
|
|
|
|
__all__ = [
|
|
"to_padded_tensor",
|
|
"as_nested_tensor",
|
|
"nested_tensor",
|
|
"narrow",
|
|
]
|
|
|
|
# Nested Tensor constructor functions
|
|
|
|
|
|
def as_nested_tensor(
|
|
tensor_list: Sequence[Tensor],
|
|
dtype: Optional[DType] = None,
|
|
device: Optional[Device] = None,
|
|
layout=None
|
|
) -> Tensor:
|
|
r"""
|
|
Constructs a nested tensor preserving autograd history from :attr:`tensor_list` a list of tensors.
|
|
|
|
.. note::
|
|
Tensors within the list are always copied by this function due to current nested tensor semantics.
|
|
|
|
Args:
|
|
tensor_list (List[Tensor]): a list of tensors with the same ndim
|
|
|
|
Keyword arguments:
|
|
dtype (:class:`torch.dtype`, optional): the desired type of returned nested tensor.
|
|
Default: if None, same :class:`torch.dtype` as leftmost tensor in the list.
|
|
device (:class:`torch.device`, optional): the desired device of returned nested tensor.
|
|
Default: if None, same :class:`torch.device` as leftmost tensor in the list
|
|
layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
|
|
Only strided and jagged layouts are supported. Default: if None, the strided 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.as_nested_tensor([a, b])
|
|
>>> nt.is_leaf
|
|
False
|
|
>>> fake_grad = torch.nested.nested_tensor([torch.ones_like(a), torch.zeros_like(b)])
|
|
>>> nt.backward(fake_grad)
|
|
>>> a.grad
|
|
tensor([1., 1., 1.])
|
|
>>> b.grad
|
|
tensor([0., 0., 0., 0., 0.])
|
|
"""
|
|
if not isinstance(tensor_list, list) or any(
|
|
not isinstance(t, Tensor) for t in tensor_list
|
|
):
|
|
raise TypeError(
|
|
"as_nested_tensor(): Expected first argument to be a list of tensors "
|
|
)
|
|
|
|
if layout is None:
|
|
layout = torch.strided
|
|
if layout == torch.strided:
|
|
return torch._nested_tensor_from_tensor_list(tensor_list, dtype, None, device, None)
|
|
elif layout == torch.jagged:
|
|
from torch.nested._internal.nested_tensor import jagged_from_list
|
|
|
|
nt, _ = jagged_from_list(tensor_list, offsets=None, device=device, dtype=dtype)
|
|
return nt
|
|
else:
|
|
raise RuntimeError(f"Specified layout is unsupported for nested tensors: {layout}")
|
|
|
|
|
|
# Note: This not only adds doc strings for the nested ops, but
|
|
# also connects the torch.nested Python namespace to the torch._C._nested builtins.
|
|
|
|
to_padded_tensor = _add_docstr(
|
|
_nested.nested_to_padded_tensor,
|
|
r"""
|
|
to_padded_tensor(input, padding, output_size=None, out=None) -> Tensor
|
|
|
|
Returns a new (non-nested) Tensor by padding the :attr:`input` nested tensor.
|
|
The leading entries will be filled with the nested data,
|
|
while the trailing entries will be padded.
|
|
|
|
.. warning::
|
|
|
|
:func:`to_padded_tensor` always copies the underlying data,
|
|
since the nested and the non-nested tensors differ in memory layout.
|
|
|
|
Args:
|
|
padding (float): The padding value for the trailing entries.
|
|
|
|
Keyword args:
|
|
output_size (Tuple[int]): The size of the output tensor.
|
|
If given, it must be large enough to contain all nested data;
|
|
else, will infer by taking the max size of each nested sub-tensor along each dimension.
|
|
out (Tensor, optional): the output tensor.
|
|
|
|
Example::
|
|
|
|
>>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))])
|
|
nested_tensor([
|
|
tensor([[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276],
|
|
[-1.9967, -1.0054, 1.8972, 0.9174, -1.4995]]),
|
|
tensor([[-1.8546, -0.7194, -0.2918, -0.1846],
|
|
[ 0.2773, 0.8793, -0.5183, -0.6447],
|
|
[ 1.8009, 1.8468, -0.9832, -1.5272]])
|
|
])
|
|
>>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0)
|
|
tensor([[[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276],
|
|
[-1.9967, -1.0054, 1.8972, 0.9174, -1.4995],
|
|
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
|
|
[[-1.8546, -0.7194, -0.2918, -0.1846, 0.0000],
|
|
[ 0.2773, 0.8793, -0.5183, -0.6447, 0.0000],
|
|
[ 1.8009, 1.8468, -0.9832, -1.5272, 0.0000]]])
|
|
>>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6))
|
|
tensor([[[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276, 1.0000],
|
|
[-1.9967, -1.0054, 1.8972, 0.9174, -1.4995, 1.0000],
|
|
[ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
|
|
[ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],
|
|
[[-1.8546, -0.7194, -0.2918, -0.1846, 1.0000, 1.0000],
|
|
[ 0.2773, 0.8793, -0.5183, -0.6447, 1.0000, 1.0000],
|
|
[ 1.8009, 1.8468, -0.9832, -1.5272, 1.0000, 1.0000],
|
|
[ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]]])
|
|
>>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2))
|
|
RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.
|
|
|
|
""",
|
|
)
|
|
|
|
def nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) -> Tensor:
|
|
r"""
|
|
Constructs a nested tensor with no autograd history (also known as a “leaf tensor”, see
|
|
:ref:`Autograd mechanics <autograd-mechanics>`) from :attr:`tensor_list` a list of tensors.
|
|
|
|
Args:
|
|
tensor_list (List[array_like]): a list of tensors, or anything that can be passed to torch.tensor,
|
|
where each element of the list has the same dimensionality.
|
|
|
|
Keyword arguments:
|
|
dtype (:class:`torch.dtype`, optional): the desired type of returned nested tensor.
|
|
Default: if None, same :class:`torch.dtype` as leftmost tensor in the list.
|
|
layout (:class:`torch.layout`, optional): the desired layout of returned nested tensor.
|
|
Only strided and jagged layouts are supported. Default: if None, the strided layout.
|
|
device (:class:`torch.device`, optional): the desired device of returned nested tensor.
|
|
Default: if None, same :class:`torch.device` as leftmost tensor in the list
|
|
requires_grad (bool, optional): If autograd should record operations on the
|
|
returned nested tensor. Default: ``False``.
|
|
pin_memory (bool, optional): If set, returned nested tensor would be allocated in
|
|
the pinned memory. Works only for CPU tensors. Default: ``False``.
|
|
|
|
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 layout is None:
|
|
layout = torch.strided
|
|
if layout == torch.strided:
|
|
return _nested.nested_tensor(
|
|
tensor_list,
|
|
dtype=dtype,
|
|
device=device,
|
|
requires_grad=requires_grad,
|
|
pin_memory=pin_memory)
|
|
elif layout == torch.jagged:
|
|
# Need to wrap lists of scalars as tensors
|
|
list_of_tensors = [t if isinstance(t, Tensor) else torch.as_tensor(t) for t in tensor_list]
|
|
|
|
from torch.nested._internal.nested_tensor import jagged_from_list
|
|
|
|
with torch.no_grad():
|
|
nt, _ = jagged_from_list(list_of_tensors, offsets=None, device=device, dtype=dtype)
|
|
|
|
nt.requires_grad_(requires_grad)
|
|
if pin_memory:
|
|
nt = nt.pin_memory() # type: ignore[assignment]
|
|
|
|
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.strided) -> Tensor:
|
|
r"""
|
|
Constructs a nested tensor (which might be a view) from :attr:`tensor`, a strided tensor. This follows
|
|
similar semantics to torch.Tensor.narrow, where in the :attr:`dim`-th dimension the new nested tensor
|
|
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]`.
|
|
|
|
There's some differences depending on the layout you use for the nested tensor. If using strided layout,
|
|
torch.narrow will do a copy of the narrowed data into a contiguous NT with strided layout, while
|
|
jagged layout narrow() will create a non-contiguous view of your original strided tensor. This particular
|
|
representation is really useful for representing kv-caches in Transformer models, as specialized
|
|
SDPA kernels can deal with format easily, resulting in performance improvements.
|
|
|
|
|
|
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 strided layout.
|
|
|
|
Example::
|
|
|
|
>>> starts = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64)
|
|
>>> lengths = torch.tensor([3, 2, 2, 1, 5], dtype=torch.int64)
|
|
>>> narrow_base = torch.randn(5, 10, 20)
|
|
>>> nt_narrowed = torch.nested.narrow(narrow_base, 1, starts, lengths, layout=torch.jagged)
|
|
>>> nt_narrowed.is_contiguous()
|
|
False
|
|
"""
|
|
if not isinstance(start, (int, SymInt, Tensor)):
|
|
raise RuntimeError("start must be an integer or a tensor")
|
|
|
|
if not isinstance(length, (int, SymInt, Tensor)):
|
|
raise RuntimeError("length must be an integer or a tensor")
|
|
|
|
if layout == torch.strided:
|
|
if isinstance(start, Tensor) or isinstance(length, Tensor):
|
|
raise RuntimeError("start and length must be integers for the strided layout NT impl")
|
|
# TODO: switch to as_nested_tensor(tensor) when it is available
|
|
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, SymInt)):
|
|
start = torch.tensor([start], device=tensor.device, dtype=torch.int64)
|
|
|
|
if isinstance(length, (int, SymInt)):
|
|
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
|