PyTorch torch.tril_indices 関数


Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.tril_indicesPyTorch で下三角行列のインデックスを生成するための関数です。下三角部分(対角線を含む)の行インデックスと列インデックスを返します。

関数定義

torch.tril_indices(row, column, offset=0, dtype=torch.long, device='cpu')

パラメータの説明:

  • row: 行数
  • column: 列数
  • offset: 対角線のオフセット量
  • dtype: 返されるデータ型
  • device: デバイス

使用例

例

import torch

# 3x3 行列の下三角インデックスを生成
row, col = torch.tril_indices(3, 3)
print("row:", row)
print("col:", col)

出力結果は次のとおりです:

row: tensor([0, 1, 1, 2, 2, 2])
col: tensor([0, 0, 1, 0, 1, 2])

例

import torch

# インデックスを生成してインデックス操作に使用
row, col = torch.tril_indices(3, 3, offset=1)

# 3x3 行列を作成
a = torch.ones(3, 3)

# インデックスを使用して下三角部分を設定
a[row, col] = 0
print(a)

出力結果は次のとおりです:

tensor([[1., 0., 0.],
        [1., 1., 0.],
        [1., 1., 1.]])

例

import torch

# 正方行列ではない場合
row, col = torch.tril_indices(3, 4)
print("row:", row)
print("col:", col)

出力結果は次のとおりです:

row: tensor([0, 1, 1, 2, 2, 2])
col: tensor([0, 0, 1, 0, 1, 2])

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張