start and end expect torch tensors

setup.py

@@ -2,7 +2,7 @@ from os import path as osp
 
 from setuptools import setup, find_packages
 
-__version__ = '0.2.4'
+__version__ = '0.2.5'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi', 'torch-unique']

test/dense_grid.json

@@ -10,7 +10,7 @@
     "name": "Start parameter",
     "position": [2, 16],
     "size": [5],
-    "start": -1,
+    "start": [-1],
     "expected": [0, 3],
     "expected_C": 4
   },
@@ -18,8 +18,8 @@
     "name": "End parameter",
     "position": [2, 16],
     "size": [5],
-    "start": 0,
-    "end": 30,
+    "start": [0],
+    "end": [30],
     "expected": [0, 3],
     "expected_C": 6
   },
@@ -43,8 +43,8 @@
     "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
     "size": [5, 5],
     "batch": [0, 0, 0, 0, 0, 1, 1],
-    "start": 0,
-    "end": 20,
+    "start": [0],
+    "end": [20],
     "expected": [0, 9, 1, 0, 4, 16, 21],
     "expected_C": 32
   }

test/sparse_grid.json

@@ -9,7 +9,7 @@
     "name": "Start parameter",
     "position": [2, 6],
     "size": [5],
-    "start": 0,
+    "start": [0],
     "expected": [0, 1]
   },
   {

test/test_dense_grid.py

@@ -20,7 +20,9 @@ def test_dense_grid_cluster_cpu(tensor, i):
     batch = data[i].get('batch')
     batch = None if batch is None else torch.LongTensor(batch)
     start = data[i].get('start')
+    start = None if start is None else torch.LongTensor(start)
     end = data[i].get('end')
+    end = None if end is None else torch.LongTensor(end)
     expected = torch.LongTensor(data[i]['expected'])
     expected_C = data[i]['expected_C']
 
@@ -37,7 +39,9 @@ def test_dense_grid_cluster_gpu(tensor, i):  # pragma: no cover
     batch = data[i].get('batch')
     batch = None if batch is None else torch.cuda.LongTensor(batch)
     start = data[i].get('start')
+    start = None if start is None else torch.cuda.LongTensor(start)
     end = data[i].get('end')
+    end = None if end is None else torch.cuda.LongTensor(end)
     expected = torch.LongTensor(data[i]['expected'])
     expected_C = data[i]['expected_C']
 

test/test_sparse_grid.py

@@ -18,16 +18,17 @@ def test_sparse_grid_cluster_cpu(tensor, i):
     position = Tensor(tensor, data[i]['position'])
     size = torch.LongTensor(data[i]['size'])
     batch = data[i].get('batch')
+    batch = None if batch is None else torch.LongTensor(batch)
     start = data[i].get('start')
+    start = None if start is None else torch.LongTensor(start)
     expected = torch.LongTensor(data[i]['expected'])
 
+    output = sparse_grid_cluster(position, size, batch, start)
+
     if batch is None:
-        output = sparse_grid_cluster(position, size, batch, start)
         assert output.tolist() == expected.tolist()
     else:
-        batch = torch.LongTensor(batch)
         expected_batch = torch.LongTensor(data[i]['expected_batch'])
-        output = sparse_grid_cluster(position, size, batch, start)
         assert output[0].tolist() == expected.tolist()
         assert output[1].tolist() == expected_batch.tolist()
 
@@ -38,15 +39,16 @@ def test_sparse_grid_cluster_gpu(tensor, i):  # pragma: no cover
     position = Tensor(tensor, data[i]['position']).cuda()
     size = torch.cuda.LongTensor(data[i]['size'])
     batch = data[i].get('batch')
+    batch = None if batch is None else torch.cuda.LongTensor(batch)
     start = data[i].get('start')
+    start = None if start is None else torch.cuda.LongTensor(start)
     expected = torch.LongTensor(data[i]['expected'])
 
+    output = sparse_grid_cluster(position, size, batch, start)
+
     if batch is None:
-        output = sparse_grid_cluster(position, size, batch, start)
         assert output.cpu().tolist() == expected.tolist()
     else:
-        batch = torch.cuda.LongTensor(batch)
         expected_batch = torch.LongTensor(data[i]['expected_batch'])
-        output = sparse_grid_cluster(position, size, batch, start)
         assert output[0].cpu().tolist() == expected.tolist()
         assert output[1].cpu().tolist() == expected_batch.tolist()

torch_cluster/__init__.py

 from .functions.grid import sparse_grid_cluster, dense_grid_cluster
 
-__version__ = '0.2.4'
+__version__ = '0.2.5'
 
 __all__ = ['sparse_grid_cluster', 'dense_grid_cluster', '__version__']

torch_cluster/functions/grid.py

@@ -22,8 +22,8 @@ def _preprocess(position, size, batch=None, start=None):
         for i in range(position.size(-1)):
             min.append(position[:, i].min())
         position = position - position.new(min)
-    elif start != 0:
-        position = position - start
+    else:
+        position = position - start.type_as(position)
 
     # If given, append batch to position tensor.
     if batch is not None:
@@ -49,6 +49,7 @@ def _fixed_cluster_size(position, size, batch=None, end=None):
     if end is None:
         return _minimal_cluster_size(position, size)
 
+    end = end.type_as(size)
     eps = 0.000001  # Simulate [start, end) interval.
     if batch is None:
         cluster_size = ((end / size).float() - eps).long() + 1