torch.nested 的源代码

```html

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",
]

# 嵌套张量构造函数


def as_nested_tensor(
    tensor_list: Sequence[Tensor],
    dtype: Optional[DType] = None,
    device: Optional[Device] = None,
    layout=None
) -> Tensor:
    r"""
    从 :attr:`tensor_list` 列表中构造一个保留自动求导历史的嵌套张量。

    .. 注意::
        由于当前嵌套张量的语义，此函数始终会复制列表中的张量。

    参数:
        tensor_list (List[Tensor]): 一个具有相同维度的张量列表

    关键字参数:
        dtype (:class:`torch.dtype`, 可选): 返回的嵌套张量的期望类型。
            默认: 如果为 None，则与列表中最左边的张量具有相同的 :class:`torch.dtype`。
        device (:class:`torch.device`, 可选): 返回的嵌套张量的期望设备。
            默认: 如果为 None，则与列表中最左边的张量具有相同的 :class:`torch.device`
        layout (:class:`torch.layout`, 可选): 返回的嵌套张量的期望布局。
            仅支持 strided 和 jagged 布局。默认: 如果为 None，则为 strided 布局。

    示例::

        >>> 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(): 期望第一个参数为张量列表 "
        )

    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"指定的布局不支持嵌套张量: {layout}")


# 注意: 这不仅为嵌套操作添加了文档字符串，
# 还将 torch.nested Python 命名空间连接到 torch._C._nested 内置函数。

to_padded_tensor = _add_docstr(
    _nested.nested_to_padded_tensor,
    r"""
to_padded_tensor(input, padding, output_size=None, out=None) -> Tensor

通过填充 :attr:`input` 嵌套张量返回一个新的 (非嵌套) 张量。
前导条目将填充嵌套数据，
而后面的条目将被填充。

.. 警告::

    :func:`to_padded_tensor` 总是复制底层数据，
    因为嵌套张量和非嵌套张量的内存布局不同。

参数:
    padding (float): 用于填充后面条目的填充值。

关键字参数:
    output_size (Tuple[int]): 输出张量的大小。
                              如果给定，它必须足够大以包含所有嵌套数据；
                              否则，将通过沿每个维度获取每个嵌套子张量的最大大小来推断。
    out (Tensor, 可选): 输出张量。

示例::

    >>> 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]],