Column-wise multiplication with matrix and vector

Is there any built-in function that multiply each column of a matrix by the corresponding element of a vector?

Example:
a = torch.tensor([[[1,2,3],[5,6,7]], [[1,3,5],[5,8,7]], [[1,1,5],[5,8,3]]]), size(2, 2, 3)
b = torch.tensor([[1,0,1], [1, 1, 0], [0, 1, 0]]), size(2, 3)
(where the first dimension is batch size)
expected output:
torch.tensor([[[1, 0, 3], [5, 0, 7]], [[1, 3, 0], [5, 8, 0]], [[0, 1, 0], [0, 8, 0]]) , size(2, 2, 3)

Thanks!

Hi xdwang!

Yes.

You can use broadcasting together with element-wise multiplication, *,
provided you use unsqueeze() so that the dimensions line up correctly:

>>> import torch
>>> torch.__version__
'1.9.0'
>>> a = torch.tensor([[[1,2,3],[5,6,7]], [[1,3,5],[5,8,7]], [[1,1,5],[5,8,3]]])
>>> b = torch.tensor([[1,0,1], [1, 1, 0], [0, 1, 0]])
>>> a * b.unsqueeze (1)
tensor([[[1, 0, 3],
         [5, 0, 7]],

        [[1, 3, 0],
         [5, 8, 0]],

        [[0, 1, 0],
         [0, 8, 0]]])

(As an aside, the sizes you give are wrong; your batch size is 3 (not 2).)

Best.

K. Frank

Thanks! It really helps!