Normalizing a multi dimensional positive tensor so that it sums to 1 along the last dimension

1

Hii,
I have a positive tensor of shape (bsz , num_heads , tgt _len , src_len). I need that it sums to 1 along the last dimension (i.e src_len dimension). I don’t want to use softmax. I want to take sum of elements along the last dimension and then divide the elements by that sum.
How should I do this? Can someone help?

Thanks in advance.

2

Try this?

X = torch.randn(bsz, num_heads, tgt_len, src_len)
s = X.sum(keepdim=True, dim=-1)
X_norm = X / s
4

Worked. Thanks. keepdim = True was important.

1 Like
5

You need keepdim=True in order for the division to broadcast correctly.