added consecutive call

test/test_random.py

@@ -7,7 +7,7 @@ def test_random():
                                    [2, 3, 6, 5, 0, 0, 4, 5, 3, 1, 3, 6, 0, 3]])
     # edge_attr = torch.Tensor([2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2])
     rid = torch.arange(edge_index.max() + 1, out=edge_index.new())
-    output = random_cluster(edge_index, rid, perm_edges=False)
+    output = random_cluster(edge_index, rid=rid, perm_edges=False)
 
-    expected_output = [0, 1, 2, 0, 4, 1, 6]
+    expected_output = [0, 1, 2, 0, 3, 1, 4]
     assert output.tolist() == expected_output

torch_cluster/functions/random.py

-from .utils import get_func
+from .utils import get_func, consecutive
 from .degree import node_degree
 from .permute import permute
 
 
-def random_cluster(edge_index, rid=None, perm_edges=True, num_nodes=None):
+def random_cluster(edge_index,
+                   batch=None,
+                   rid=None,
+                   perm_edges=True,
+                   num_nodes=None):
+
     num_nodes = edge_index.max() + 1 if num_nodes is None else num_nodes
     row, col = permute(edge_index, num_nodes, rid, perm_edges)
     degree = node_degree(row, num_nodes, out=row.new())
@@ -12,4 +17,10 @@ def random_cluster(edge_index, rid=None, perm_edges=True, num_nodes=None):
     func = get_func('random', cluster)
     func(cluster, row, col, degree)
 
-    return cluster
+    cluster, u = consecutive(cluster)
+
+    if batch is None:
+        return cluster
+    else:
+        # TODO: Fix
+        return cluster, batch

torch_cluster/functions/utils.py

@@ -26,7 +26,7 @@ def get_type(max, cuda):
         return torch.cuda.LongTensor if cuda else torch.LongTensor
 
 
-def consecutive(tensor, return_batch=None):
+def consecutive(tensor):
     size = tensor.size()
     u = unique(tensor.view(-1))
     len = u[-1] + 1