Element wise batch matrix multiplication of a row with every other row in matrix, in PyTorch
I have two matrices of sizes (30, 24, 512) respectively where 30 is the batch size. Let us call them A and B.
Now what I need to do is this:
For every batch in A, I want to compute element-wise batch matrix multiplication of each row in a single batch of A with each row in a single batch of B and sum them.
In other words, for every batch, I have a (24, 512) matrix on left-hand side (A) and on right-hand side (B). Now for each row (1 * 512) in A, I want to compute element-wise matrix multiplication of that row with each of the (24 * 512) rows in B, one by one, and sum them.
This operation would thus result in a (1 * 512) sized vector.
Done for each row in A, it would result in a (24 * 512) sized matrix and then, done for each batch it would result in (30, 24, 512) sized matrix, which is the preferred dimension of my output.
How do I go about doing this?
I do not want to use for loops, since that would be inefficient. I also do not want to use repeat method for the same reason.

I figured out the answer to my own question. Instead of computing the element-wise product of each row and then wanting to compute the sum, I summed the rows and then took the element-wise product, which made the problem (and the solution) much simpler.
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

