Fix: error on instance from batch_x missing from batch_y in nearest()

test/test_nearest.py

@@ -33,3 +33,38 @@ def test_nearest(dtype, device):
 
     out = nearest(x, y)
     assert out.tolist() == [0, 0, 1, 1, 2, 2, 3, 3]
+
+    # Invalid input: instance 1 only in batch_x
+    batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
+    batch_y = tensor([0, 0, 0, 0], torch.long, device)
+    with pytest.raises(ValueError):
+        out = nearest(x, y, batch_x, batch_y)
+
+    # Invalid input: instance 1 only in batch_x (implicitly as batch_y=None)
+    with pytest.raises(ValueError):
+        out = nearest(x, y, batch_x, batch_y=None)
+
+    # Valid input: instance 1 only in batch_y
+    batch_x = tensor([0, 0, 0, 0, 0, 0, 0, 0], torch.long, device)
+    batch_y = tensor([0, 0, 1, 1], torch.long, device)
+    out = nearest(x, y, batch_x, batch_y)
+    assert out.tolist() == [0, 0, 1, 1, 0, 0, 1, 1]
+
+    # Invalid input: instance 2 only in batch_x
+    # (i.e.instance in the middle missing)
+    batch_x = tensor([0, 0, 1, 1, 2, 2, 3, 3], torch.long, 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()
@@ -67,10 +72,31 @@ def nearest(x: torch.Tensor, y: torch.Tensor,
         else:
             ptr_y = torch.tensor([0, y.size(0)], device=y.device)
 
+        # if an instance in batch_x is non-empty, it must be
+        # non-empty in batch_y as well
+        instance_nonempty_x = (ptr_x[:-1] != ptr_x[1:])
+        instance_nonempty_y = (ptr_y[:len(ptr_x)-1] != ptr_y[1:len(ptr_x)])
+        if (len(ptr_x) > len(ptr_y) or
+                (instance_nonempty_x & ~instance_nonempty_y).any()):
+            raise ValueError("Some batch index occurs in batch_x "
+                             "that does not occur in batch_y")
+
         return torch.ops.torch_cluster.nearest(x, y, ptr_x, ptr_y)
     else:
+        if (batch_x is None) != (batch_y is None):
+            raise ValueError("Either both or none of batch_x, batch_y "
+                             "may be None")
+
         # Translate and rescale x and y to [0, 1].
         if batch_x is not None and batch_y is not None:
+            # if an instance in batch_x is non-empty, it must be
+            # non-empty in batch_y as well
+            if not torch.isin(torch.unique_consecutive(batch_x),
+                              torch.unique_consecutive(batch_y),
+                              assume_unique=True).all():
+                raise ValueError("Some batch index occurs in batch_x "
+                                 "that does not occur in batch_y")
+
             assert x.dim() == 2 and batch_x.dim() == 1
             assert y.dim() == 2 and batch_y.dim() == 1
             assert x.size(0) == batch_x.size(0)