torch.triu_indices

torch.triu_indices(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) → 张量

返回一个2×N张量中矩阵的上三角部分的索引，其中第一行包含所有索引的行坐标，第二行包含列坐标。索引按行和列排序。

矩阵的上三角部分定义为对角线及其上方的元素。

参数 offset 控制要考虑的对角线。如果 offset = 0，则保留主对角线及以上的所有元素。正值会排除主对角线上方的对角线，同样，负值会包含主对角线下方的对角线。主对角线是索引集 $\lbrace (i, i) \rbrace$ 对于 $i \in [0, \min\{d_{1}, d_{2}\} - 1]$ 其中 $d_{1}, d_{2}$ 是矩阵的维度。

注意

当在CUDA上运行时，row * col 必须小于 $2^{59}$ 以防止在计算过程中发生溢出。

Parameters

• 行 (int) – 二维矩阵中的行数。

• col (int) – 二维矩阵中的列数。

• offset (int) – 相对于主对角线的对角偏移量。 默认值：如果未提供，则为0。

Keyword Arguments

• dtype (torch.dtype, 可选) – 返回张量所需的数据类型。 默认值：如果 None，则为 torch.long。

• 设备 (torch.device, 可选) – 返回张量所需的设备。 默认值：如果 None，则使用默认张量类型的当前设备 （参见 torch.set_default_device()）。device 将是 CPU 用于 CPU 张量类型，以及当前 CUDA 设备用于 CUDA 张量类型。

• 布局 (torch.layout, 可选) – 目前仅支持 torch.strided。

示例：

>>> a = torch.triu_indices(3, 3)
>>> a
tensor([[0, 0, 0, 1, 1, 2],
        [0, 1, 2, 1, 2, 2]])

>>> a = torch.triu_indices(4, 3, -1)
>>> a
tensor([[0, 0, 0, 1, 1, 1, 2, 2, 3],
        [0, 1, 2, 0, 1, 2, 1, 2, 2]])

>>> a = torch.triu_indices(4, 3, 1)
>>> a
tensor([[0, 0, 1],
        [1, 2, 2]])