torch.diag() 取矩阵对角线元素，torch.diag_embed() 指定值变成对角矩阵

1、torch.diag()

import torch
a = torch.randn(3, 3)
print(a)
tensor([[ 0.7594,  0.8073, -0.1344],
        [-1.7335, -0.4356, -0.0055],
        [ 1.8326,  0.3900, -0.9933]])
diag = torch.diag(a)  # 取 a 对角线元素，输出为 1*3
print(diag)
tensor([ 0.7594, -0.4356, -0.9933])

2、torch.diag_embed()

import torch

tensor([ 0.7594, -0.4356, -0.9933])
a_diag = torch.diag_embed(diag)  # 由 diag 变为三维 3*3
tensor([[ 0.7594,  0.0000,  0.0000],
        [ 0.0000, -0.4356,  0.0000],
        [ 0.0000,  0.0000, -0.9933]])