better consecutive impl

test/utils/test_consecutive.py

@@ -4,17 +4,11 @@ from torch_cluster.functions.utils.consecutive import consecutive
 
 
 def test_consecutive_cpu():
-    vec = torch.LongTensor([0, 2, 3])
-    assert consecutive(vec).tolist() == [0, 1, 2]
-
-    vec = torch.LongTensor([0, 3, 2, 2, 3])
-    assert consecutive(vec).tolist() == [0, 2, 1, 1, 2]
+    x = torch.LongTensor([0, 3, 2, 2, 3])
+    assert consecutive(x).tolist() == [0, 2, 1, 1, 2]
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
 def test_consecutive_gpu():  # pragma: no cover
-    vec = torch.cuda.LongTensor([0, 2, 3])
-    assert consecutive(vec).cpu().tolist() == [0, 1, 2]
-
-    vec = torch.cuda.LongTensor([0, 3, 2, 2, 3])
-    assert consecutive(vec).cpu().tolist() == [0, 2, 1, 1, 2]
+    x = torch.cuda.LongTensor([0, 3, 2, 2, 3])
+    assert consecutive(x).cpu().tolist() == [0, 2, 1, 1, 2]

torch_cluster/functions/utils/consecutive.py

@@ -14,13 +14,21 @@ def _get_type(max_value, cuda):
 
 
 def consecutive(x):
-    size = x.size()
+    initial_size = x.size()
+
+    # Compute unique vector.
     u = unique(x.view(-1))
-    len = u[-1] + 1
-    max = u.size(0)
-    type = _get_type(max, x.is_cuda)
-    arg = type(len)
-    arg[u] = torch.arange(0, max, out=type(max))
-    x = arg[x.view(-1)]
-    x = x.view(size).long()
+
+    # Compute mask with mask[u[0]] = 0, mask[u[1]] = 1, ...
+    # As mask can get very big (dependent on the maximum value in `x`, we want
+    # to take the least possible amount of space on disk (`_get_type`).
+    max_value = u[-1] + 1
+    mask = _get_type(u.size(0), x.is_cuda)(max_value)
+    mask[u] = torch.arange(0, u.size(0), out=mask.new())
+
+    # Select the values in `mask` based on `x` and reshape to initial size.
+    x = mask[x.view(-1)]
+    x = x.view(initial_size)
+    x = x.long()
+
     return x

torch_cluster/functions/utils/degree.py

@@ -4,5 +4,5 @@ import torch
 def node_degree(index, num_nodes, out=None):
     out = index.new(num_nodes) if out is None else out
     zero = torch.zeros(num_nodes, out=out)
-    one = torch.ones(index.size(0), out=zero.new(index.size(0)))
+    one = torch.ones(index.size(0), out=zero.new())
     return zero.scatter_add_(0, index, one)