Scalar operation for n-dimensional tensor

Hi all

I have a question regarding scalar operation on n-dimensional tensors.

I have two tensors shaped as follows;

x torch.Size([30, 48, 201, 229])
gamma torch.Size([30, 48])

I would like to add/multiply the gamma to each element of matrix [201, 229] in the tensor x for each 48 dimension for 30 batches.

How can I achieve this?
Thank you so much for your help in advance!

This is a case for broadcasting: You add size-1 dimensions to gamma so it has ones for the last two axes: x * gamma[:, :, None, None] the : takes a full dimension from the original gamma, the None adds a singleton dimension which is then treated with broadcasting.
You could also pretend to be fancy by using einsum: torch.einsum('ijkl,ij->ijkl', x, gamma) will give the same result. If you remove letters from the last bit (following the ->) of the equation string, einsum will sum over these.

Best regards


1 Like

Great! It works! Thank you so much!