Repository: pyg-team/pytorch_geometric
Pull request #6875 (2023-03-08 09:30:10)

Changed the interface of `sort_edge_index` to only return single `edge_index` information in case the `edge_attr` argument is not specified

codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6875?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6875](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6875?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (c766eb5) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/97c1a532324578d69ec924f3616cc482df024fd4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (97c1a53) will **decrease** coverage by `0.01%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head c766eb5 differs from pull request most recent head deb8efe. Consider uploading reports for the commit deb8efe to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #6875      +/-   ##
==========================================
- Coverage   91.52%   91.52%   -0.01%     
==========================================
  Files         427      427              
  Lines       23160    23159       -1     
==========================================
- Hits        21198    21197       -1     
  Misses       1962     1962              
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6875?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/transforms/to\_sparse\_tensor.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6875?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvdG9fc3BhcnNlX3RlbnNvci5weQ==) | `97.14% <100.00%> (ø)` | |
| [torch\_geometric/utils/sort\_edge\_index.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6875?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NvcnRfZWRnZV9pbmRleC5weQ==) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index da50a7b42..3eaebf87a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -98,6 +98,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Changed the interface of `sort_edge_index` to only return single `edge_index` information in case the `edge_attr` argument is not specified ([#6875](https://github.com/pyg-team/pytorch_geometric/issues/6875))
 - Fixed a bug in `to_hetero` when using an uninitialized submodule without implementing `reset_parameters` ([#6863](https://github.com/pyg-team/pytorch_geometric/issues/6790))
 - Fixed a bug in `get_mesh_laplacian` ([#6790](https://github.com/pyg-team/pytorch_geometric/issues/6790))
 - Fixed a bug in which masks were not properly masked in `GNNExplainer` on link prediction tasks ([#6787](https://github.com/pyg-team/pytorch_geometric/pull/6787))
diff --git a/torch_geometric/transforms/to_sparse_tensor.py b/torch_geometric/transforms/to_sparse_tensor.py
index 65b5fc7dc..030daa918 100644
--- a/torch_geometric/transforms/to_sparse_tensor.py
+++ b/torch_geometric/transforms/to_sparse_tensor.py
@@ -56,9 +56,11 @@ class ToSparseTensor(BaseTransform):
                     keys.append(key)
                     values.append(value)
 
-            store.edge_index, values = sort_edge_index(store.edge_index,
-                                                       values,
-                                                       sort_by_row=False)
+            store.edge_index, values = sort_edge_index(
+                store.edge_index,
+                values,
+                sort_by_row=False,
+            )
 
             for key, value in zip(keys, values):
                 store[key] = value
diff --git a/torch_geometric/utils/sort_edge_index.py b/torch_geometric/utils/sort_edge_index.py
index 3b07bfd76..63b2f2b36 100644
--- a/torch_geometric/utils/sort_edge_index.py
+++ b/torch_geometric/utils/sort_edge_index.py
@@ -7,31 +7,30 @@ from torch_geometric.utils import index_sort
 
 from .num_nodes import maybe_num_nodes
 
+MISSING = '???'
+
 
 @torch.jit._overload
-def sort_edge_index(edge_index, edge_attr=None, num_nodes=None,
-                    sort_by_row=True):
-    # type: (Tensor, Optional[bool], Optional[int], bool) -> Tensor  # noqa
+def sort_edge_index(edge_index, edge_attr, num_nodes, sort_by_row):
+    # type: (Tensor, str, Optional[int], bool) -> Tensor  # noqa
     pass
 
 
 @torch.jit._overload
-def sort_edge_index(edge_index, edge_attr=None, num_nodes=None,
-                    sort_by_row=True):
-    # type: (Tensor, Tensor, Optional[int], bool) -> Tuple[Tensor, Tensor]  # noqa
+def sort_edge_index(edge_index, edge_attr, num_nodes, sort_by_row):
+    # type: (Tensor, Optional[Tensor], Optional[int], bool) -> Tuple[Tensor, Optional[Tensor]]  # noqa
     pass
 
 
 @torch.jit._overload
-def sort_edge_index(edge_index, edge_attr=None, num_nodes=None,
-                    sort_by_row=True):
+def sort_edge_index(edge_index, edge_attr, num_nodes, sort_by_row):
     # type: (Tensor, List[Tensor], Optional[int], bool) -> Tuple[Tensor, List[Tensor]]  # noqa
     pass
 
 
 def sort_edge_index(
     edge_index: Tensor,
-    edge_attr: Union[Optional[Tensor], List[Tensor]] = None,
+    edge_attr: Union[Optional[Tensor], List[Tensor], str] = MISSING,
     num_nodes: Optional[int] = None,
     sort_by_row: bool = True,
 ) -> Union[Tensor, Tuple[Tensor, Tensor], Tuple[Tensor, List[Tensor]]]:
@@ -48,8 +47,8 @@ def sort_edge_index(
         sort_by_row (bool, optional): If set to :obj:`False`, will sort
             :obj:`edge_index` column-wise.
 
-    :rtype: :class:`LongTensor` if :attr:`edge_attr` is :obj:`None`, else
-        (:class:`LongTensor`, :obj:`Tensor` or :obj:`List[Tensor]]`)
+    :rtype: :class:`LongTensor` if :attr:`edge_attr` is not passed, else
+        (:class:`LongTensor`, :obj:`Optional[Tensor]` or :obj:`List[Tensor]]`)
 
     Examples:
 
@@ -77,9 +76,11 @@ def sort_edge_index(
 
     edge_index = edge_index[:, perm]
 
+    if edge_attr is None:
+        return edge_index, None
     if isinstance(edge_attr, Tensor):
         return edge_index, edge_attr[perm]
     elif isinstance(edge_attr, (list, tuple)):
         return edge_index, [e[perm] for e in edge_attr]
-    else:
-        return edge_index
+
+    return edge_index


Test patch:
diff --git a/test/utils/test_convert.py b/test/utils/test_convert.py
index 7da4e3f4c..cd9d6ea32 100644
--- a/test/utils/test_convert.py
+++ b/test/utils/test_convert.py
@@ -433,11 +433,8 @@ def test_to_cugraph(edge_weight, directed, relabel_nodes):
         cu_edge_weight = edge_list['weights'].to_pandas().values
         cu_edge_weight = torch.from_numpy(cu_edge_weight)
 
-    if edge_weight is None:
-        cu_edge_index = sort_edge_index(cu_edge_index)
-    else:
-        cu_edge_index, cu_edge_weight = sort_edge_index(
-            cu_edge_index, cu_edge_weight)
+    cu_edge_index, cu_edge_weight = sort_edge_index(cu_edge_index,
+                                                    cu_edge_weight)
 
     assert torch.equal(edge_index, cu_edge_index)
     if edge_weight is not None:
@@ -476,11 +473,8 @@ def test_from_cugraph(edge_weight, directed, relabel_nodes):
     )
 
     cu_edge_index, cu_edge_weight = from_cugraph(G)
-    if edge_weight is None:
-        cu_edge_index = sort_edge_index(cu_edge_index)
-    else:
-        cu_edge_index, cu_edge_weight = sort_edge_index(
-            cu_edge_index, cu_edge_weight)
+    cu_edge_index, cu_edge_weight = sort_edge_index(cu_edge_index,
+                                                    cu_edge_weight)
 
     assert torch.equal(edge_index, cu_edge_index)
     if edge_weight is not None:
diff --git a/test/utils/test_sort_edge_index.py b/test/utils/test_sort_edge_index.py
index 8575e7971..1e57fbfc3 100644
--- a/test/utils/test_sort_edge_index.py
+++ b/test/utils/test_sort_edge_index.py
@@ -1,4 +1,7 @@
+from typing import List, Optional, Tuple
+
 import torch
+from torch import Tensor
 
 from torch_geometric.utils import sort_edge_index
 
@@ -10,6 +13,50 @@ def test_sort_edge_index():
     out = sort_edge_index(edge_index)
     assert out.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
 
+    out = sort_edge_index(edge_index, None)
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1] is None
+
     out = sort_edge_index(edge_index, edge_attr)
     assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
     assert out[1].tolist() == [[4], [3], [2], [1]]
+
+
+def test_sort_edge_index_jit():
+    @torch.jit.script
+    def wrapper1(edge_index: Tensor) -> Tensor:
+        return sort_edge_index(edge_index)
+
+    @torch.jit.script
+    def wrapper2(
+        edge_index: Tensor,
+        edge_attr: Optional[Tensor],
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        return sort_edge_index(edge_index, edge_attr)
+
+    @torch.jit.script
+    def wrapper3(
+        edge_index: Tensor,
+        edge_attr: List[Tensor],
+    ) -> Tuple[Tensor, List[Tensor]]:
+        return sort_edge_index(edge_index, edge_attr)
+
+    edge_index = torch.tensor([[2, 1, 1, 0], [1, 2, 0, 1]])
+    edge_attr = torch.tensor([[1], [2], [3], [4]])
+
+    out = wrapper1(edge_index)
+    assert out.size() == edge_index.size()
+
+    out = wrapper2(edge_index, None)
+    assert out[0].size() == edge_index.size()
+    assert out[1] is None
+
+    out = wrapper2(edge_index, edge_attr)
+    assert out[0].size() == edge_index.size()
+    assert out[1].size() == edge_attr.size()
+
+    out = wrapper3(edge_index, [edge_attr, edge_attr])
+    assert out[0].size() == edge_index.size()
+    assert len(out[1]) == 2
+    assert out[1][0].size() == edge_attr.size()
+    assert out[1][1].size() == edge_attr.size()