Perform a batch matrix - multiple weight matrices multiplications in pytorch
I have a batch of matrices A with size torch.Size([batch_size, 9, 5]) and weight matrices B with size torch.Size([3, 5, 6]). In Keras, a simple K.dot(A, B) is able to handle the matrix multiplication to give an output with size (batch_size, 9, 3, 6). Here, each row in A is multiplied to the 3 matrices in B to form a (3x6) matrix.

How do you perform a similar operation in torch. From the documentation, torch.bmm requires that A and B must have the same batch size, so I tried this:

B = B.unsqueeze(0).repeat((batch_size, 1, 1, 1))
B.size() # torch.Size([batch_size, 3, 5, 6])
torch.bmm(A,B) # gives an error



  RuntimeError: invalid argument 2: expected 3D tensor, got 4D


Well, the error is expected but how do I perform such an operation? 

You can use einstein notation to describe the operation you want as bxy,iyk-&gt;bxik. So, you can use einsum to calculate it.

torch.einsum('bxy,iyk-&gt;bxik', (A, B)) will give you the answer you want.
