Update utils.py

Author: cszhangzhen

utils.py

@@ -10,12 +10,31 @@
 from torch_geometric.transforms import OneHotDegree
 from torch_geometric.data import Data, InMemoryDataset
 from torch_geometric.utils import add_remaining_self_loops, remove_self_loops, dense_to_sparse
-from torch_geometric.utils import softmax, degree, sort_edge_index
+from torch_geometric.utils import softmax, degree
 from torch_scatter import scatter
 from torch_cluster import random_walk
 from torch_sparse import spspmm, coalesce
 
 
+def sort_edge_index(edge_index, edge_attr=None, num_nodes=None):
+    r"""Row-wise sorts edge indices :obj:`edge_index`.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        edge_attr (Tensor, optional): Edge weights or multi-dimensional
+            edge features. (default: :obj:`None`)
+        num_nodes (int, optional): The number of nodes, *i.e.*
+            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
+
+    :rtype: (:class:`LongTensor`, :class:`Tensor`)
+    """
+
+    idx = edge_index[0] * num_nodes + edge_index[1]
+    perm = idx.argsort()
+
+    return edge_index[:, perm], None if edge_attr is None else edge_attr[perm]
+
+
 class BinaryFuncDataset(InMemoryDataset):
     def __init__(self, root, name, transform=None, pre_transform=None, pre_filter=None):
         self.dir_name = os.path.join(root, name)