grid cpu

cpu/grid.cpp

@@ -2,18 +2,17 @@
 
 at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
                 at::Tensor end) {
-  size = size.toType(pos.type());
-  start = start.toType(pos.type());
-  end = end.toType(pos.type());
-
   pos = pos - start.view({1, -1});
-  auto num_voxels = ((end - start) / size).toType(at::kLong);
-  num_voxels = (num_voxels + 1).cumsum(0);
-  num_voxels -= num_voxels.data<int64_t>()[0];
-  num_voxels.data<int64_t>()[0] = 1;
 
-  auto cluster = pos / size.view({1, -1});
-  cluster = cluster.toType(at::kLong);
+  auto num_voxels = ((end - start) / size).toType(at::kLong) + 1;
+  num_voxels = num_voxels.cumprod(0);
+
+  num_voxels = at::cat({at::ones(1, num_voxels.options()), num_voxels}, 0);
+  auto index = empty(size.size(0), num_voxels.options());
+  arange_out(index, size.size(0));
+  num_voxels = num_voxels.index_select(0, index);
+
+  auto cluster = (pos / size.view({1, -1})).toType(at::kLong);
   cluster *= num_voxels.view({1, -1});
   cluster = cluster.sum(1);
 

setup.cfg

@@ -5,4 +5,4 @@ description-file = README.md
 test = pytest
 
 [tool:pytest]
-addopts = --capture=no --cov
+addopts = --capture=no

torch_cluster/grid.py

@@ -28,9 +28,7 @@ def grid_cluster(pos, size, start=None, end=None):
     start = pos.t().min(dim=1)[0] if start is None else start
     end = pos.t().max(dim=1)[0] if end is None else end
 
-    if pos.is_cuda:
-        cluster = grid_cuda.grid(pos, size, start, end)
-    else:
-        cluster = grid_cpu.grid(pos, size, start, end)
+    op = grid_cuda.grid if pos.is_cuda else grid_cpu.grid
+    cluster = op(pos, size, start, end)
 
     return cluster