PyTorch torch.diag 関数

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

torch.diagPyTorch において、対角行列を作成したり、テンソルの対角要素を抽出したりするための関数です。

関数定義

torch.diag(input, diagonal=0, out=None)

パラメータの説明：

  • input: 入力テンソル
  • diagonal: 対角線のインデックス。0 は主対角線、正の値は上側の対角線、負の値は下側の対角線を表します。

使用例

例

import torch

# 1次元テンソルを作成
x = torch.tensor([1, 2, 3])

# 対角行列を作成
y = torch.diag(x)
print(y)

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

tensor([[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]])

例

import torch

# 行列から対角要素を抽出
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 主対角線を抽出
y = torch.diag(x)
print(y)

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

tensor([1, 5, 9])

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

その他の拡張