sample doesn't work correctly when nodes are repeated in input

https://github.com/rusty1s/pytorch_sparse/blob/3ec1eacf1c1cbf22c2fc95cff6ddd1ccc5d2efa6/csrc/cpu/neighbor_sample_cpu.cpp#L108

Working on https://github.com/pyg-team/pytorch_geometric/issues/4026 I discovered that this sampling function does not work properly when the input_node array has duplicates. In this case the number of samples can be > number of edges and the index i doesn’t do what you’d expect it to.

Another strange behaviour is when using the function as

torch.ops.torch_sparse.neighbor_sample(
            torch.Tensor([0, 1, 2, 3]).type(torch.long),
            torch.Tensor([1, 0, 1]).type(torch.long),
            torch.Tensor([0, 2]).type(torch.long),
            [1],
            False,
            False,
        )

The returned values are

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

but the (row, col) combination (2, 0) and (0, 2) don’t make sense, they’re not part of the original adjacency matrix.

I’m happy to dive in a bit deeper and figure out a fix (for at least the first issue, the second issue I am not sure if its actually a problem or not yet), but thought I’d check first if this is known or maybe my understanding is incorrect.

[Padarn]

ohh I see, the row/col are relative to the nodes in the output? If so, then agreed.

And if this is the case… then actually the first issue is also a non-issue. Sorry for the false alarm!

[Padarn]

Closing this one as we’ve agreed not to do anything about it in pytorch_sparse, thanks @rusty1s