Remove negative indexing

Author: nicolas-chaulet

cpu/src/ball_query.cpp

@@ -84,9 +84,13 @@ std::pair<at::Tensor, at::Tensor> batch_ball_query(at::Tensor support, at::Tenso
     auto options_dist = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
 
     int max_count = 0;
-    auto batch_access = query_batch.accessor<int64_t, 1>();
+    auto q_batch_access = query_batch.accessor<int64_t, 1>();
+    auto s_batch_access = support_batch.accessor<int64_t, 1>();
+
+    auto batch_size = q_batch_access[query_batch.size(0) - 1] + 1;
+    TORCH_CHECK(batch_size == (s_batch_access[support_batch.size(0) - 1] + 1),
+                "Both batches need to have the same number of samples.")
 
-    auto batch_size = batch_access[-1] + 1;
     query_batch = degree(query_batch, batch_size);
     query_batch = at::cat({at::zeros(1, query_batch.options()), query_batch.cumsum(0)}, 0);
     support_batch = degree(support_batch, batch_size);

test/test_ballquerry.py

@@ -1,11 +1,16 @@
 import unittest
 import torch
-from torch_points_kernels import ball_query
 import numpy.testing as npt
 import numpy as np
 from sklearn.neighbors import KDTree
+import os
+import sys
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
 
-from . import run_if_cuda
+from test import run_if_cuda
+from torch_points_kernels import ball_query
 
 
 class TestBall(unittest.TestCase):
@@ -76,10 +81,10 @@ def test_simple_gpu(self):
         npt.assert_array_almost_equal(dist2, dist2_answer)
 
     def test_simple_cpu(self):
-        x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0], [0.1, 0, 0]]).to(torch.float)
+        x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0], [10.1, 0, 0]]).to(torch.float)
         y = torch.tensor([[0, 0, 0]]).to(torch.float)
 
-        batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long()
+        batch_x = torch.from_numpy(np.asarray([0, 0, 0, 0])).long()
         batch_y = torch.from_numpy(np.asarray([0])).long()
 
         idx, dist2 = ball_query(1.0, 2, x, y, mode="PARTIAL_DENSE", batch_x=batch_x, batch_y=batch_y)
@@ -93,6 +98,17 @@ def test_simple_cpu(self):
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)
 
+    
+    def test_breaks(self):
+        x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0], [10.1, 0, 0]]).to(torch.float)
+        y = torch.tensor([[0, 0, 0]]).to(torch.float)
+
+        batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long()
+        batch_y = torch.from_numpy(np.asarray([0])).long()
+        
+        with self.assertRaises(RuntimeError):
+            idx, dist2 = ball_query(1.0, 2, x, y, mode="PARTIAL_DENSE", batch_x=batch_x, batch_y=batch_y)
+
     def test_random_cpu(self):
         a = torch.randn(100, 3).to(torch.float)
         b = torch.randn(50, 3).to(torch.float)