如何在每批中创建具有不同元素的火炬对角矩阵？

Question

我想创建一个张量像

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

即，给定大小为（1，n）的火炬张量B时，我要创建大小为（n，3,3）的火炬张量A，以使A [i]为B [i] * （大小为3x3的恒等矩阵）。

不使用'for句子'，如何创建它？

Answer 1

使用torch.einsum（爱因斯坦的总和与积的符号）

A = torch.eye(3)
b = torch.tensor([1.0, 2.0, 3.0])
torch.einsum('ij,k->kij', A, b)

将返回：

tensor([[[1., 0., 0.],
     [0., 1., 0.],
     [0., 0., 1.]],

    [[2., 0., 0.],
     [0., 2., 0.],
     [0., 0., 2.]],

    [[3., 0., 0.],
     [0., 3., 0.],
     [0., 0., 3.]]])