torch - 如何在 PyTorch 中构造一个 3D 张量，其中每个 2D 子张量都是对角矩阵？

Question

假设我有 2D 张量，index_in_batch * diag_ele。我怎样才能得到一个3D张量index_in_batch * Matrix(谁是对角矩阵，由drag_ele构造)？

torch.diag()仅当输入为一维时构造对角矩阵，当输入为二维时返回对角元素。

Answer 1

import torch

a = torch.rand(2, 3)
print(a)
b = torch.eye(a.size(1))
c = a.unsqueeze(2).expand(*a.size(), a.size(1))
d = c * b
print(d)

输出

 0.5938  0.5769  0.0555
 0.9629  0.5343  0.2576
[torch.FloatTensor of size 2x3]


(0 ,.,.) = 
  0.5938  0.0000  0.0000
  0.0000  0.5769  0.0000
  0.0000  0.0000  0.0555

(1 ,.,.) = 
  0.9629  0.0000  0.0000
  0.0000  0.5343  0.0000
  0.0000  0.0000  0.2576
[torch.FloatTensor of size 2x3x3]