torch_geometric.nn.nearest argument order flipped

[gegewen]

🐛 Describe the bug

The variable order in nearest is flipped.

The documentation stated "Finds for each element in y the k nearest point in x", which means the output dimension should be the same as x. However, when I used the following code to test this function, the output dimension is the same as x.

import torch
from torch_geometric.nn import nearest

x = torch.tensor([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
y = torch.tensor([[-1.0, 0.0], [1.0, 0.0]])
cluster = nearest(x, y)
print(cluster, cluster.shape)

output:

tensor([0, 0, 1, 1]) torch.Size([4])

The issue was fixed after I flipped x and y.

I think this has to do with the fact that torch_geometric.nn.nearest directly adopted the ordering convention form torch_cluster.nearest(), which has a flipped order fromtorch_cluster.knn().

Versions

[pip3] numpy==2.0.2
[pip3] torch==2.4.1
[pip3] torch_cluster==1.6.3
[pip3] torch_geometric==2.5.3
[pip3] torch_scatter==2.1.2
[pip3] torchaudio==2.4.1
[pip3] torchvision==0.19.1

[AyushmanGarg]

@akihironitta Can I be assigned to this bug?