PyTorch torch.diag_embed 関数

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

torch.diag_embedは、PyTorch において入力テンソルを対角要素として新しいテンソルに埋め込むための関数です。入力の最後の次元に対角線を作成します。

関数定義

torch.diag_embed(input, diagonal=0, offset=0)

引数の説明：

  • input: 入力テンソル
  • diagonal/offset: 対角インデックス

使用例

例

import torch

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

# 対角埋め込みを作成
y = torch.diag_embed(x)
print(y)

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

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

例

import torch

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

# 対角埋め込みを作成
y = torch.diag_embed(x)
print(y.shape)
print(y)

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

torch.Size([2, 2, 2])
tensor([[[1, 0],
         [0, 2]],

        [[3, 0],
         [0, 4]]])

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

その他の拡張