recursive

csrc/cpu/metis_wrapper_cpu.cpp

@@ -4,8 +4,8 @@
 
 #include "utils.h"
 
-torch::Tensor partition_kway_cpu(torch::Tensor rowptr, torch::Tensor col,
-                                 int64_t num_parts) {
+torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                            int64_t num_parts, bool recursive) {
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
 
@@ -17,8 +17,13 @@ torch::Tensor partition_kway_cpu(torch::Tensor rowptr, torch::Tensor col,
   auto part = torch::empty(nvtxs, rowptr.options());
   auto part_data = part.data_ptr<int64_t>();
 
-  METIS_PartGraphKway(&nvtxs, &ncon, xadj, adjncy, NULL, NULL, NULL, &num_parts,
-                      NULL, NULL, NULL, &objval, part_data);
+  if (recursive) {
+    METIS PartGraphRecursive(&nvtxs, &ncon, xadj, adjncy, NULL, NULL, NULL,
+                             &num_parts, NULL, NULL, NULL, &objval, part_data);
+  } else {
+    METIS_PartGraphKway(&nvtxs, &ncon, xadj, adjncy, NULL, NULL, NULL,
+                        &num_parts, NULL, NULL, NULL, &objval, part_data);
+  }
 
   return part;
 }

csrc/cpu/metis_wrapper_cpu.h

@@ -2,5 +2,5 @@
 
 #include <torch/extension.h>
 
-torch::Tensor partition_kway_cpu(torch::Tensor rowptr, torch::Tensor col,
-                                 int64_t num_parts);
+torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                            int64_t num_parts, bool recursive);

csrc/metis_wrapper.cpp

@@ -9,8 +9,8 @@
 PyMODINIT_FUNC PyInit__metis_wrapper(void) { return NULL; }
 #endif
 
-torch::Tensor partition_kway(torch::Tensor rowptr, torch::Tensor col,
-                             int64_t num_parts) {
+torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
+                        int64_t num_parts, bool recursive) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
     AT_ERROR("No CUDA version supported");
@@ -18,9 +18,9 @@ torch::Tensor partition_kway(torch::Tensor rowptr, torch::Tensor col,
     AT_ERROR("Not compiled with CUDA support");
 #endif
   } else {
-    return partition_kway_cpu(rowptr, col, num_parts);
+    return partition_kway_cpu(rowptr, col, num_parts, recursive);
   }
 }
 
-static auto registry = torch::RegisterOperators().op(
-    "torch_sparse::partition_kway", &partition_kway);
+static auto registry =
+    torch::RegisterOperators().op("torch_sparse::partition", &partition);

torch_sparse/metis.py

@@ -5,12 +5,13 @@ from torch_sparse.tensor import SparseTensor
 from torch_sparse.permute import permute
 
 
-def partition_kway(
-        src: SparseTensor,
-        num_parts: int) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
+def partition(
+    src: SparseTensor, num_parts: int, recursive: bool = False
+) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
 
     rowptr, col = src.storage.rowptr().cpu(), src.storage.col().cpu()
-    cluster = torch.ops.torch_sparse.partition_kway(rowptr, col, num_parts)
+    cluster = torch.ops.torch_sparse.partition(rowptr, col, num_parts,
+                                               recursive)
     cluster = cluster.to(src.device())
 
     cluster, perm = cluster.sort()
@@ -20,5 +21,4 @@ def partition_kway(
     return out, partptr, perm
 
 
-SparseTensor.partition_kway = lambda self, num_parts: partition_kway(
-    self, num_parts)
+SparseTensor.partition = lambda self, num_parts: partition(self, num_parts)