torch.triu_indices

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

返回一個 row x col 矩陣的上三角部分的索引，以一個 2xN 的 Tensor 返回，其中第一行包含所有索引的行座標，第二行包含列座標。索引按行然後按列排序。

矩陣的上三角部分定義為對角線上的元素及其上方的元素。

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, optional) – 返回張量的期望資料型別，僅支援 torch.int、torch.long。預設值：如果為 None，則為 torch.long。

• device (torch.device, 可選) – 返回張量的所需裝置。預設：如果為 None，則使用當前裝置作為預設張量型別（參見 torch.set_default_device()）。對於 CPU 張量型別，device 將是 CPU；對於 CUDA 張量型別，device 將是當前的 CUDA 裝置。

• layout (torch.layout, optional) – 目前僅支援 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]])