CODEWITHSUNDEEP
pytorch
gtoques
gtoques

Reputation: 233

Multiplying PyTorch tensors of different shape

I have a torch tensor of shape (32, 100, 50) and another of shape (32,100). Call these A and B respectively. I want to element-wise multiply A and B, such that each of the 50 elements at A[i, j, :] get multiplied by B[i, j], i.e.like multiplying a vector with a scalar. How can I do this via broadcasting rules?

Upvotes: 1

Views: 1084

Answers (1)

jhso
jhso

Reputation: 3283

Just add a singleton dimension to the second tensor, for example:

a = torch.randn([32,100,50])
b = torch.randint(10,[32,100])
b = b[:,:,None] #or .unsqueeze(-1)
c = a * b 
assert (c[0,0,0]/a[0,0,0]).int() == b[0,0] and (c[0,0,1]/a[0,0,1]).int() == b[0,0]

The assert on the end is just to prove that adjacent elements in the last dimension is multiplied by the same element of b.

Upvotes: 1

Related Questions