如何构建一个3D Tensor,其中每个2D子张量都是PyTorch中的对角矩阵？

Question

考虑我有2D Tensor , index_in_batch * diag_ele. 如何获得3D Tensor index_in_batch * Matrix(谁是对角矩阵,由drag_ele构造)？

在torch.diag()仅当输入是一维,并返回对角元素时输入是2D构建体对角矩阵.

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]

Answer 2

使用torch.diag_embed：

>>> a = torch.randn(2, 3)
>>> torch.diag_embed(a)
tensor([[[ 1.5410,  0.0000,  0.0000],
         [ 0.0000, -0.2934,  0.0000],
         [ 0.0000,  0.0000, -2.1788]],

        [[ 0.5684,  0.0000,  0.0000],
         [ 0.0000, -1.0845,  0.0000],
         [ 0.0000,  0.0000, -1.3986]]])