How to implement 4D tensor multiplication?

Given A: [B, N, K, K], B: [B, S, K, K].
The shape of expected matrix multiplication result: [B, N, S, K, K].
How to use torch.matmul to get this result?


Assuming that you want to reduce dimension -1 of A and dimension -2 of B,
you can do the following so that the batching works fine:

torch.matmul(A.unsqueeze(3), B.unsqueeze(2))


I have tried your solution. But I met some errors. I use the code below.

a = torch.rand(2, 8, 3, 3)
b = torch.rand(2, 4, 3, 3)
ans = torch.matmul(a.unsqueeze(3), b.unsqueeze(2))

And I got the error:

ans = torch.matmul(a.unsqueeze(3), b.unsqueeze(2))
RuntimeError: The size of tensor a (8) must match the size of tensor b (4) at non-singleton dimension 1

I know the problem. The following code works.

torch.matmul(A.unsqueeze(2), B.unsqueeze(1))

I’m curious about how it works. Could you please explain the reason?


This function may be help you
It multiplys two matrixs of any dimension.

Yeah, I have tried torch.einsum and got the same results.

torch.einsum('abik, ackj -> abcij', A, B)


Ho my bad I miscounted the dimensions.

What the unsqueeze does is to make the sizes 2, 1, 8, 3, 3 and 2, 4, 1, 3, 3. So that matmul can broadcast on these two dimensions of size 1 and do the matrix product you want.

Got it. Thanks for your explanation.