Поэлементное пакетное матричное умножение строки на каждую другую строку в матрице в PyTorch

Question

У меня есть две матрицы размеров (30, 24, 512) соответственно, где 30 - размер партии. Давайте назовем их A и B.

Теперь мне нужно сделать следующее:

Для каждого пакета в A я хочу вычислить поэлементную матрицу пакетов умножение каждой строки в одной партии из A на каждую строку в одной партии из B и их суммирование. Другими словами, для каждого пакета у меня есть матрица (24, 512) слева (A) и справа (B). Теперь для каждой строки (1 * 512) в A я хочу вычислить поэлементное матричное умножение этой строки на каждую из (24 * 512) строк в B, одну за другой, и просуммировать их. Таким образом, эта операция приведет к вектору размером (1 * 512). Выполнено для каждой строки в A, это приведет к матрице размера (24 * 512), а затем, выполненной для каждого пакета, это приведет к матрице размера (30, 24, 512), что является предпочтительным размером моего вывода.

Как мне go об этом? Я не хочу использовать циклы for, поскольку это было бы неэффективно. Я также не хочу использовать метод repeat по той же причине.

Answer 1

Я придумал ответ на свой вопрос. Вместо того, чтобы вычислять поэлементное произведение каждой строки, а затем желать вычислить сумму, я суммировал строки, а затем взял поэлементное произведение, что значительно упростило задачу (и решение).

import torch

a = torch.randn(30, 24, 512)
b = torch.randn(30, 24, 512)

# Step 1: Summing b along dimension 1 and unsqueezing
# Dimension of b_sum is 30*1*512
b_sum = torch.sum(b, dim=1).unsqueeze(1)

# Step 2: Element-wise product 
result = a * b_sum