bugfix, added tests

test/test_grid.py

+import pytest
+import torch
+from torch_cluster import grid_cluster
+
+from .utils import tensors, Tensor
+
+
+@pytest.mark.parametrize('tensor', tensors)
+def test_grid_cluster_cpu(tensor):
+    position = Tensor(tensor, [[0, 0], [9, 9], [2, 8], [2, 2], [8, 3]])
+    size = torch.LongTensor([5, 5])
+    expected = torch.LongTensor([0, 3, 1, 0, 2])
+
+    output = grid_cluster(position, size)
+    assert output.tolist() == expected.tolist()
+
+    output = grid_cluster(position.expand(2, 5, 2), size)
+    assert output.tolist() == expected.expand(2, 5).tolist()
+
+    expected = torch.LongTensor([0, 1, 3, 2, 4])
+    batch = torch.LongTensor([0, 0, 1, 1, 1])
+    output = grid_cluster(position, size, batch)
+    assert output.tolist() == expected.tolist()
+
+    output = grid_cluster(position.expand(2, 5, 2), size, batch.expand(2, 5))
+    assert output.tolist() == expected.expand(2, 5).tolist()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
+@pytest.mark.parametrize('tensor', tensors)
+def test_grid_cluster_gpu(tensor):  # pragma: no cover
+    pass

torch_cluster/functions/grid.py

@@ -29,7 +29,7 @@ def grid_cluster(position, size, batch=None):
     max = position.max(dim=0)[0]
     while max.dim() > 1:
         max = max.max(dim=0)[0]
-    c_max = torch.ceil(max / size.type_as(max)).long()
+    c_max = torch.floor(max.double() / size.double() + 1).long()
     c_max = torch.clamp(c_max, min=1)
     C = c_max.prod()
 

torch_cluster/functions/utils.py

@@ -2,6 +2,7 @@ import torch
 from torch_unique import unique
 
 from .._ext import ffi
+print(ffi.__dict__)
 
 
 def get_func(name, tensor):

torch_cluster/src/generic/cuda.c

@@ -3,7 +3,7 @@
 #else
 
 void cluster_(grid)(int C, THCudaLongTensor *output, THCTensor *position, THCTensor *size, THCudaLongTensor *count) {
-  return cluster_kernel_(grid)(state, C, output, position, size, count);
+  /* return cluster_kernel_(grid)(state, C, output, position, size, count); */
 }
 
 #endif