Improve the numerical stability of cosine similarity

[chenzhekl]

Related: pytorch/pytorch#31378
Current implementation:

pytorch_cluster/csrc/cuda/knn_cuda.cu

Line 56 in dbcafbe

tmp_dist = Cosine<scalar_t>::dot(x, y, n_x, n_y, dim) /

As introduced in the linked PR for PyTorch, normalizing the vectors first and then doing the inner product could prevent overflow for vectors with large norms. In addition, this could ensure the calculated cosine similarity is within the range [-1, +1].

[rusty1s]

This is clever. Do you want to contribute this feature? :)

[chenzhekl]

Sure, I would love to.

[rusty1s]

Thanks!