knn example

[ldv1]

Hi,

Using the knn example from pool.knn, i.e.

import torch
from torch_geometric.nn import knn

x = torch.tensor([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.tensor([[-1.0, 0.0], [1.0, 0.0]])
batch_y = torch.tensor([0, 0])
assign_index = knn(x, y, 2, batch_x, batch_y)

I got

tensor([[0, 0, 1, 1],
        [0, 1, 2, 3]])

though I was expecting

tensor([[0, 1],
        [2, 3]])

Do I err?

[songsong0425]

What's the meaning of your expectation?
To my knowledge, knn returns the 2 nearest neighbors from y to x when [0, 0, 1, 1] means the node with the y label and [0, 1, 2, 3] means the x nodes. You can simply check the process by drawing the nodes on the coordinate.

[ldv1]

I see, my interpretation of the output was wrong. Thanks.