add one optional arg to 'partition'

test/test_metis.py

@@ -8,12 +8,16 @@ from .utils import devices
 @pytest.mark.parametrize('device', devices)
 def test_metis(device):
     weighted_mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
-    mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=False)
+    mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=False, sort_strategy=True)
+    assert partptr.numel() == 3
+    assert perm.numel() == 6
+
+    mat, partptr, perm = weighted_mat.partition(num_parts=2, recursive=False, sort_strategy=False)
     assert partptr.numel() == 3
     assert perm.numel() == 6
 
     unweighted_mat = SparseTensor.from_dense(torch.ones((6, 6), device=device))
-    mat, partptr, perm = unweighted_mat.partition(num_parts=2, recursive=True)
+    mat, partptr, perm = unweighted_mat.partition(num_parts=2, recursive=True, sort_strategy=True)
     assert partptr.numel() == 3
     assert perm.numel() == 6
 

torch_sparse/metis.py

@@ -11,27 +11,46 @@ def cartesian1d(x, y):
     return coos.split(1, dim=1)
 
 
-def metis_weight(x):
+def metis_weight1(x):
+    sorted_x = x.sort()[0]
+    diff = sorted_x[1:] - sorted_x[:-1]
+    if diff.sum() == 0:
+        return None
+    xmin, xmax = sorted_x[[0, -1]]
+    srange = xmax - xmin
+    min_diff = diff.min()
+    scale = (min_diff / srange).item()
+    tick, arange = scale.as_integer_ratio()
+    x_ratio = (x - xmin) / srange
+    return (x_ratio * arange + tick).long()
+
+
+def metis_weight2(x):
     t1, t2 = cartesian1d(x, x)
     diff = t1 - t2
     diff = diff[diff != 0]
     if len(diff) == 0:
         return None
-    res = diff.abs().min()
-    bod = x.max() - x.min()
-    scale = (res / bod).item()
+    xmin, xmax = x.min(), x.max()
+    srange = xmax - xmin
+    min_diff = diff.abs().min()
+    scale = (min_diff / srange).item()
     tick, arange = scale.as_integer_ratio()
-    x_ratio = (x - x.min()) / bod
+    x_ratio = (x - xmin) / srange
     return (x_ratio * arange + tick).long()
 
 
-def partition(src: SparseTensor, num_parts: int, recursive: bool = False
+def metis_weight(x, sort_strategy=True):
+    return metis_weight1(x) if sort_strategy else metis_weight2(x)
+
+
+def partition(src: SparseTensor, num_parts: int, recursive: bool = False, sort_strategy=True,
               ) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
     rowptr, col, value = src.csr()
     rowptr, col = rowptr.cpu(), col.cpu()
     if value is not None and value.dim() == 1:
         value = value.detach().cpu()
-        value = metis_weight(value)
+        value = metis_weight(value, sort_strategy)
     cluster = torch.ops.torch_sparse.partition(rowptr, col, value, num_parts,
                                                recursive)
     cluster = cluster.to(src.device())