Multiply a 3d and a 2d tensor

1

I have a 3d tensor A = [5, 12, 12] and a 2d tensor B = [5, 12]

So basically 5 2d tensors and 5 1d tensors

And I want to multiply them to get a 2d tensor C = [5, 12]
i.e. I want to multiply every 1d tensor with every 2d tensor to get a 1d tensor and then stack them to get a 2d tensor.

I did something that I thought would work, but I am not getting the desired result.

def mul(A: torch.tensor, B: torch.tensor) -> torch.tensor:
    t_res = []
    for i in range(len(A):
        t_res.append(B[i].matmul(Ai]))
    
    return torch.stack(tuple(t_res), dim = 0)

Is there a better/easier way to do this?

2

unsquezeing B should work:

A = torch.randn(5, 12, 12)
B = torch.randn(5, 12)

t_res = []
for i in range(len(A)):
    t_res.append(B[i].matmul(A[i]))
ref = torch.stack(tuple(t_res), dim = 0)

out = torch.matmul(B.unsqueeze(1), A)

print((out.squeeze(1) - ref).abs().max())
# tensor(0.)
1 Like