[FIXED] PyTorch: Computing the norm of batched tensors

Issue

I have tensor t with shape (Batch_Size x Dims) and another tensor v with shape (Vocab_Size x Dims). I'd like to produce a tensor d with shape (Batch_Size x Vocab_Size), such that d[i,j] = norm(t[i] - v[j]).

Doing this for a single tensor (no batches) is trivial: d = torch.norm(v - t), since t would be broadcast. How can I do this when the tensors have batches?

Solution

Insert unitary dimensions into v and t to make them (1 x Vocab_Size x Dims) and (Batch_Size x 1 x Dims) respectively. Next, take the broadcasted difference to get a tensor of shape (Batch_Size x Vocab_Size x Dims). Pass that to torch.norm along with the optional dim=2 argument so that the norm is taken along the last dimension. This will result in the desired (Batch_Size x Vocab_Size) tensor of norms.

d = torch.norm(v.unsqueeze(0) - t.unsqueeze(1), dim=2)

Edit: As pointed out by @KonstantinosKokos in the comments, due to the broadcasting rules used by numpy and pytorch, the leading unitary dimension on v does not need to be explicit. I.e. you can use

d = torch.norm(v - t.unsqueeze(1), dim=2)

Answered By - jodag