torch.tril_indices

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

返回一个 2xN 的 Tensor，其中包含 row 行 col 列矩阵的下三角部分的索引。该 2xN Tensor 的第一行包含所有索引的行坐标，第二行包含列坐标。索引按行然后按列排序。

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

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

注意

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

参数

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

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

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

关键字参数

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

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

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

示例

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

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

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