torch.triu_indices

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

返回一个2xN的张量，其中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}$ 才能防止计算过程中溢出。

参数

• row (int) – 2D矩阵的行数。

• col (int) – 2D矩阵的列数。

• offset (int) – 对角线与主对角线的偏移量。默认为0。

关键字参数

• dtype (torch.dtype, 可选) – 返回张量的期望数据类型，仅支持 torch.int, torch.long。如果未指定，默认为 torch.long。

• device (torch.device, 可选) – 返回张量的期望设备。如果未指定，则使用默认张量类型的当前设备（请参阅 torch.set_default_device()）。对于CPU张量类型，device 将是CPU；对于CUDA张量类型，device 将是当前CUDA设备。

• layout (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]])