Graph partition based on balance_edge (rusty1s#309)

* add balance_edge for graph partition * add Tensor defition in metis.py * update --------- Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>

Author: ZhengHongming888

torch_sparse/metis.py

@@ -1,11 +1,13 @@
-from typing import Tuple, Optional
+from typing import Optional, Tuple
 
 import torch
-from torch_sparse.tensor import SparseTensor
+from torch import Tensor
+
 from torch_sparse.permute import permute
+from torch_sparse.tensor import SparseTensor
 
 
-def weight2metis(weight: torch.Tensor) -> Optional[torch.Tensor]:
+def weight2metis(weight: Tensor) -> Optional[Tensor]:
  sorted_weight = weight.sort()[0]
  diff = sorted_weight[1:] - sorted_weight[:-1]
  if diff.sum() == 0:
@@ -20,16 +22,24 @@ def weight2metis(weight: torch.Tensor) -> Optional[torch.Tensor]:
 
 
 def partition(
- src: SparseTensor, num_parts: int, recursive: bool = False,
- weighted: bool = False, node_weight: Optional[torch.Tensor] = None
-) -> Tuple[SparseTensor, torch.Tensor, torch.Tensor]:
+ src: SparseTensor,
+ num_parts: int,
+ recursive: bool = False,
+ weighted: bool = False,
+ node_weight: Optional[Tensor] = None,
+ balance_edge: bool = False,
+) -> Tuple[SparseTensor, Tensor, Tensor]:
 
  assert num_parts >= 1
  if num_parts == 1:
  partptr = torch.tensor([0, src.size(0)], device=src.device())
  perm = torch.arange(src.size(0), device=src.device())
  return src, partptr, perm
 
+ if balance_edge and node_weight:
+ raise ValueError("Cannot set 'balance_edge' and 'node_weight' at the "
+ "same time in 'torch_sparse.partition'")
+
  rowptr, col, value = src.csr()
  rowptr, col = rowptr.cpu(), col.cpu()
 
@@ -41,6 +51,10 @@ def partition(
  else:
  value = None
 
+ if balance_edge:
+ node_weight = col.new_zeros(col.size(0))
+ node_weight.scatter_add_(0, col, torch.ones_like(col))
+
  if node_weight is not None:
  assert node_weight.numel() == rowptr.numel() - 1
  node_weight = node_weight.view(-1).detach().cpu()