В PyTorch / Numpy, как умножить строки матрицы на «матрицы» в трехмерном тензоре?

Question

Например, a = torch.Tensor([[1,2],[3,4]]) (для numpy это просто a = np.array([[1,2],[3,4]])) и b = torch.ones((2,2,2)),

Я хотел бы скомпоновать каждую строку a с двумя 2x2 матрицами и получить новую матрицу [[3,3],[7,7]] (т.е. [1,2]*[[1,1],[1,1]]=[3,3], [3,4]*[[1,1],[1,1]]=[7,7]). Возможно ли этого добиться? Спасибо!

Answer 1

Я считаю это уродливым решением, но, возможно, это то, чего вы хотите достичь:

a = torch.Tensor([[1,2],[3,4]])
b = torch.ones((2,2,2))

A = torch.mm(a[0].view(-1, 2), b[0])
B = torch.mm(a[1].view(-1, 2), b[1])
res = torch.cat([A, B], dim=0)
print(res)

выход:

tensor([[3., 3.],
        [7., 7.]])

Answer 2

Вы можете просто использовать обычное матричное умножение a @ b:

a = np.array([[1,2],[3,4]])
b = np.ones((2,2,2))

print(a @ b)

Вывод:

[[[3. 3.]
  [7. 7.]]

 [[3. 3.]
  [7. 7.]]]

Каждая «строка» результата (т.е. первый индекс) будетотдельная матрица 2 x 2:

print((a @ b)[1])

Выход:

[[3. 3.]
 [7. 7.]]