Element wise batch matrix multiplication of a row with every other row in matrix, in PyTorch

Question

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.

remorax · Accepted Answer · 2020-07-14 10:12:53Z

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

Please mark your answer with green tick as solved for future readers, thanks for participating! You can read about it here
@SzymonMaszke oh yes I did try to earlier, but I get an error saying I can only accept my answer tomorrow.

Collectives™ on Stack Overflow

Element wise batch matrix multiplication of a row with every other row in matrix, in PyTorch

1 Answer 1

3 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Related