How to multiply each element of a vector with every element of another vector using PyTorch

I am trying to multiply two tensors such that every element of each row vector is multiplied with every element of the corresponding row vector of the other tensor.

Say, I have

t1 = [ [a1, a2, a3], [b1, b2, b3], [c1, c2, c3] ] #Shape 23
t2 = [ [x1, x2, x3], [y1, y2, y3], [z1, z2, z3] ] #Shape 2
3

I expect the output tensor to be :

out_t = [ [a1x1, a1x2, a1x3, a2x1, a2x2, a2x3, a3x1, a3x2, a3x3], [b1y1, b1y2, b1y3, b2y1, b2y2, b2y3, b3y1,b3y2,b3y3]] #Shape 2*9

Is there a clean way of achieving this is PyTorch?

I’d recommend to look at broadcasting in the docs:

out_t = t1[:, :, None]*t2[:, None, :]

Indexing with None (just like numpy.newaxis) inserts a single axis, like unsqueeze(dim) would, too, but for tasks like this, I like the explicit way to see exactly how it aligns (similarly, the trailing colon for t2 is just to make the dimensions explicit).
There also is torch.einsum if you want to be fance (or for more elaborate use-cases), but here, I’d use the above in my own code.

Best regards

Thomas

1 Like