1

2

3

4

5

import torch

data = torch.tensor([1,2,3,4])

data_two = torch.diag(data,0)

print(data_two)

1

2

3

4

5

import torch

data = torch.tensor(float('inf')).cuda().repeat(3)

data_two = torch.diag(data,0)

print(data_two)

1

2

3

4

5

6

7

8

9

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])

1

2

3

4

5

6

7

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]])