Unsqueeze a certain number of times depending on a condition

1

Hi,
I want to create a function that takes a batch of inputs X and a set of coeffs C and then multiply each element of the batch X[i] by the corresponding coeff C[i]. One has X.shape = (B, N1, ..., Nn) and C.shape = (B).
To do so, I want to unsqueeze C so that it is broadcastable with X and then return C * X. However I want my function to accept any shape while X.size(0) == C.size(0) so I cannot predict how many times I must unsqueeze the last dimension of C. What would be the best way to do so?

2

You don’t need to call unsqueeze() so many times.
Assuming X and C are like below

X = torch.rand(2, 3, 4, 5) # [B, N1, N2, N3]
C = torch.rand(2) # [B]
B = C.size(0)

C = C.unsqueeze(1) # [B, 1]
C_expand = C.expand_as(X.view(B, -1)) # [B, N1 x N2 x N3]
C_new = C.view_as(X) # [B, N1, N2, N3]

Now then you can do what you want

3

Oh I see the trick now, thanks a lot! I was so obsessed by this unsqueeze that I forgot I could just expand my tensor without additional memory usage…