rename

csrc/cpu/metis_cpu.cpp

@@ -12,34 +12,36 @@
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
                             torch::optional<torch::Tensor> optional_value,
-                            torch::optional<torch::Tensor> vweights,
+                            torch::optional<torch::Tensor> optional_node_weight,
                             int64_t num_parts, bool recursive) {
 #ifdef WITH_METIS
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
+
   if (optional_value.has_value()) {
     CHECK_CPU(optional_value.value());
     CHECK_INPUT(optional_value.value().dim() == 1);
     CHECK_INPUT(optional_value.value().numel() == col.numel());
   }
 
-  if (vweights.has_value()) {
-    CHECK_CPU(vweights.value());
-    CHECK_INPUT(vweights.value().dim() == 1);
-    CHECK_INPUT(vweights.value().numel() == rowptr.numel() - 1);
+  if (optional_node_weight.has_value()) {
+    CHECK_CPU(optional_node_weight.value());
+    CHECK_INPUT(optional_node_weight.value().dim() == 1);
+    CHECK_INPUT(optional_node_weight.value().numel() == rowptr.numel() - 1);
   }
 
   int64_t nvtxs = rowptr.numel() - 1;
   int64_t ncon = 1;
   auto *xadj = rowptr.data_ptr<int64_t>();
   auto *adjncy = col.data_ptr<int64_t>();
+
   int64_t *adjwgt = NULL;
   if (optional_value.has_value())
     adjwgt = optional_value.value().data_ptr<int64_t>();
 
   int64_t *vwgt = NULL;
-  if (vweights.has_value())
-    vwgt = vweights.value().data_ptr<int64_t>();
+  if (optional_node_weight.has_value())
+    vwgt = optional_node_weight.value().data_ptr<int64_t>();
 
   int64_t objval = -1;
   auto part = torch::empty(nvtxs, rowptr.options());
@@ -62,11 +64,11 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
 // needs mt-metis installed via:
 // ./configure --shared --edges64bit --vertices64bit --weights64bit
 //             --partitions64bit
-torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                               torch::optional<torch::Tensor> optional_value,
-                               torch::optional<torch::Tensor> vweights,
-                               int64_t num_parts, bool recursive,
-                               int64_t num_workers) {
+torch::Tensor
+mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                 torch::optional<torch::Tensor> optional_value,
+                 torch::optional<torch::Tensor> optional_node_weight,
+                 int64_t num_parts, bool recursive, int64_t num_workers) {
 #ifdef WITH_MTMETIS
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
@@ -76,10 +78,10 @@ torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
     CHECK_INPUT(optional_value.value().numel() == col.numel());
   }
 
-  if (vweights.has_value()) {
-    CHECK_CPU(vweights.value());
-    CHECK_INPUT(vweights.value().dim() == 1);
-    CHECK_INPUT(vweights.value().numel() == rowptr.numel() - 1);
+  if (optional_node_weight.has_value()) {
+    CHECK_CPU(optional_node_weight.value());
+    CHECK_INPUT(optional_node_weight.value().dim() == 1);
+    CHECK_INPUT(optional_node_weight.value().numel() == rowptr.numel() - 1);
   }
 
   mtmetis_vtx_type nvtxs = rowptr.numel() - 1;
@@ -87,11 +89,14 @@ torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
   mtmetis_adj_type *xadj = (mtmetis_adj_type *)rowptr.data_ptr<int64_t>();
   mtmetis_vtx_type *adjncy = (mtmetis_vtx_type *)col.data_ptr<int64_t>();
   mtmetis_wgt_type *adjwgt = NULL;
+
   if (optional_value.has_value())
     adjwgt = optional_value.value().data_ptr<int64_t>();
+
   mtmetis_wgt_type *vwgt = NULL;
-  if (vweights.has_value())
-    vwgt = vweights.value().data_ptr<int64_t>();
+  if (optional_node_weight.has_value())
+    vwgt = optional_node_weight.value().data_ptr<int64_t>();
+
   mtmetis_pid_type nparts = num_parts;
   mtmetis_wgt_type objval = -1;
   auto part = torch::empty(nvtxs, rowptr.options());

csrc/cpu/metis_cpu.h

@@ -4,11 +4,11 @@
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
                             torch::optional<torch::Tensor> optional_value,
-                            torch::optional<torch::Tensor> vweights,
+                            torch::optional<torch::Tensor> optional_node_weight,
                             int64_t num_parts, bool recursive);
 
-torch::Tensor mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                               torch::optional<torch::Tensor> optional_value,
-                               torch::optional<torch::Tensor> vweights,
-                               int64_t num_parts, bool recursive,
-                               int64_t num_workers);
+torch::Tensor
+mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                 torch::optional<torch::Tensor> optional_value,
+                 torch::optional<torch::Tensor> optional_node_weight,
+                 int64_t num_parts, bool recursive, int64_t num_workers);

csrc/metis.cpp

@@ -13,7 +13,7 @@ PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; }
 
 torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
                         torch::optional<torch::Tensor> optional_value,
-                        torch::optional<torch::Tensor> vweights,
+                        torch::optional<torch::Tensor> optional_node_weight,
                         int64_t num_parts, bool recursive) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
@@ -22,13 +22,14 @@ torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
     AT_ERROR("Not compiled with CUDA support");
 #endif
   } else {
-    return partition_cpu(rowptr, col, optional_value, vweights, num_parts, recursive);
+    return partition_cpu(rowptr, col, optional_value, optional_node_weight,
+                         num_parts, recursive);
   }
 }
 
 torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
                            torch::optional<torch::Tensor> optional_value,
-                           torch::optional<torch::Tensor> vweights,
+                           torch::optional<torch::Tensor> optional_node_weight,
                            int64_t num_parts, bool recursive,
                            int64_t num_workers) {
   if (rowptr.device().is_cuda()) {
@@ -38,8 +39,8 @@ torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
     AT_ERROR("Not compiled with CUDA support");
 #endif
   } else {
-    return mt_partition_cpu(rowptr, col, optional_value, vweights, num_parts, recursive,
-                            num_workers);
+    return mt_partition_cpu(rowptr, col, optional_value, optional_node_weight,
+                            num_parts, recursive, num_workers);
   }
 }
 

test/test_metis.py

 import pytest
+from itertools import product
+
 import torch
 from torch_sparse.tensor import SparseTensor
 
@@ -12,32 +14,24 @@ except RuntimeError:
 
 
 @pytest.mark.skipif(not with_metis, reason='Not compiled with METIS support')
-@pytest.mark.parametrize('device', devices)
-def test_metis(device):
-    value1 = torch.randn(6 * 6, device=device).view(6, 6)
-    value2 = torch.arange(6 * 6, dtype=torch.long, device=device).view(6, 6)
-    value3 = torch.ones(6 * 6, device=device).view(6, 6)
-
-    vwgts = torch.rand(6, device=device)
-
-    for value in [value1, value2, value3]:
-        for vwgt in [None, vwgts]:
-            mat = SparseTensor.from_dense(value)
-
-            _, partptr, perm = mat.partition(num_parts=2, recursive=False,
-                                             vweights=vwgt,
-                                             weighted=True)
-            assert partptr.numel() == 3
-            assert perm.numel() == 6
-
-            _, partptr, perm = mat.partition(num_parts=2, recursive=False,
-                                             vweights=vwgt,
-                                             weighted=False)
-            assert partptr.numel() == 3
-            assert perm.numel() == 6
-
-            _, partptr, perm = mat.partition(num_parts=1, recursive=False,
-                                             vweights=vwgt,
-                                             weighted=True)
-            assert partptr.numel() == 2
-            assert perm.numel() == 6
+@pytest.mark.parametrize('device,weighted', product(devices, [False, True]))
+def test_metis(device, weighted):
+    mat1 = torch.randn(6 * 6, device=device).view(6, 6)
+    mat2 = torch.arange(6 * 6, dtype=torch.long, device=device).view(6, 6)
+    mat3 = torch.ones(6 * 6, device=device).view(6, 6)
+
+    vec1 = None
+    vec2 = torch.rand(6, device=device)
+
+    for mat, vec in product([mat1, mat2, mat3], [vec1, vec2]):
+        mat = SparseTensor.from_dense(mat)
+
+        _, partptr, perm = mat.partition(num_parts=1, recursive=False,
+                                         weighted=weighted, node_weight=vec)
+        assert partptr.numel() == 2
+        assert perm.numel() == 6
+
+        _, partptr, perm = mat.partition(num_parts=2, recursive=False,
+                                         weighted=weighted, node_weight=vec)
+        assert partptr.numel() == 3
+        assert perm.numel() == 6

torch_sparse/metis.py

@@ -20,9 +20,8 @@ def weight2metis(weight: torch.Tensor) -> Optional[torch.Tensor]:
 
 
 def partition(
-        src: SparseTensor, num_parts: int, recursive: bool = False,
-        vweights: torch.tensor = None,
-        weighted: bool = False
+    src: SparseTensor, num_parts: int, recursive: bool = False,
+    weighted: bool = False, node_weight: torch.tensor = None
 ) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
 
     assert num_parts >= 1
@@ -42,14 +41,14 @@ def partition(
     else:
         value = None
 
-    if vweights is not None:
-        assert vweights.numel() == rowptr.numel() - 1
-        vweights = vweights.view(-1).detach().cpu()
-        if vweights.is_floating_point():
-            vweights = weight2metis(vweights)
+    if node_weight is not None:
+        assert node_weight.numel() == rowptr.numel() - 1
+        node_weight = node_weight.view(-1).detach().cpu()
+        if node_weight.is_floating_point():
+            node_weight = weight2metis(node_weight)
 
-    cluster = torch.ops.torch_sparse.partition(rowptr, col, value, vweights, num_parts,
-                                               recursive)
+    cluster = torch.ops.torch_sparse.partition(rowptr, col, value, node_weight,
+                                               num_parts, recursive)
     cluster = cluster.to(src.device())
 
     cluster, perm = cluster.sort()