Error on unsorted batch_x/batch_y in nearest()

Author: Vuenc

test/test_nearest.py

@@ -56,3 +56,15 @@ def test_nearest(dtype, device):
     batch_y = tensor([0, 1, 3, 3], torch.long, device)
     with pytest.raises(ValueError):
         out = nearest(x, y, batch_x, batch_y)
+
+    # Invalid input: batch_x unsorted
+    batch_x = tensor([0, 0, 1, 0, 0, 0, 0], torch.long, device)
+    batch_y = tensor([0, 0, 1, 1], torch.long, device)
+    with pytest.raises(ValueError):
+        out = nearest(x, y, batch_x, batch_y)
+
+    # Invalid input: batch_y unsorted
+    batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
+    batch_y = tensor([0, 0, 1, 0], torch.long, device)
+    with pytest.raises(ValueError):
+        out = nearest(x, y, batch_x, batch_y)

torch_cluster/nearest.py

@@ -42,6 +42,11 @@ def nearest(x: torch.Tensor, y: torch.Tensor,
     y = y.view(-1, 1) if y.dim() == 1 else y
     assert x.size(1) == y.size(1)
 
+    if batch_x is not None and (batch_x[1:] - batch_x[:-1] < 0).any():
+        raise ValueError("batch_x is not sorted")
+    if batch_y is not None and (batch_y[1:] - batch_y[:-1] < 0).any():
+        raise ValueError("batch_y is not sorted")
+
     if x.is_cuda:
         if batch_x is not None:
             assert x.size(0) == batch_x.numel()