Tensors multiplication 3D * 2D

I have a 3D tensor A = [4 x 10 x 128] and a 2D tensor B = [ 4 x 10]. I want to multiply them and have a 3D tensor C = [1 x 10 x 128].

Currently I am trying this:

result = torch.mm(A.view(-1, 4), B)
result = result.view(1, 10, 128)

But I get an error that the size [1 x 10 x 128] is invalid for input with 12800 elements.

Any suggestions?

You could do a batch matrix multiply (I’m not sure if this is what you’re looking for?) by turning the 128 dimension into the batch dimension:

A = A.permute(2, 1, 0)  # A is now 128 x 10 x 4
A.bmm(B)
1 Like

I tried that but I get the following error :

RuntimeError: dimension out of range ( expected to be in range of [-2, 1] but got 2.

Any ideas what this is?

Oops, my bad. The right function to use is matmul:

import torch
A = torch.randn(4, 10, 128)
B = torch.randn(4, 10)
C = A.permute(2, 1, 0).matmul(B)

You won’t be able to get a tensor C of size [1 x 10 x 128] by a matrix multiplication. As the error says, if

A.view(-1, 4) # is of size [1280 x 4] and
B # is of size [4 x 10]

then,

result = torch.mm(A.view(-1, 4), B) # is of size [1280 x 10] = [12800]

which cannot be rewritten as a tensor of size [1 x 10 x 128] = [1280].

Sometimes a small recap of the basics helps!