Search code examples
pytorchcosine-similarity

Computing Cosine Distance with Differently shaped tensors

I have the following tensor representing a word vector

A = (2, 500)

Where the first dimension is the BATCH dimension (i.e. A contains two word vectors each with 500 elements)

I also have the following tensor

B = (10, 500)

I want to compute the cosine distance between A and B such that I get

C = (2, 10, 1)

i.e for each row in A compute the cosine distance with each row in B

I looked at using torch.nn.functional.F.cosine_similarity however this doesn't work as the dimensions must be the same.

Whats the best efficient way of achieving this in pytorch?

Solution

  • Use broadcasting technique with unsqueeze

    import torch.nn.functional as F

C = F.cosine_similarity(A.unsqueeze(1), B, dim=-1)

print(C.shape)
# torch.size([2,10])