Repository: pyg-team/pytorch_geometric
Pull request #7814 (2023-07-30 04:00:00)

Let `coalesce()` and `sort_edge_index()` accept `edge_index` as a tuple of tensors
This PR allows `coalesce()` and `sort_edge_index()` to accept `edge_index` as a tensor or a tuple of tensors.

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

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

```diff
@@            Coverage Diff             @@
##           master    #7814      +/-   ##
==========================================
- Coverage   91.93%   91.57%   -0.37%     
==========================================
  Files         455      455              
  Lines       25659    25659              
==========================================
- Hits        23590    23497      -93     
- Misses       2069     2162      +93     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7814?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/coalesce.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7814?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2NvYWxlc2NlLnB5) | `100.00% <100.00%> (ø)` | |

... and [21 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7814/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

: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)

EdisonLeeeee: Thank you for the PR. However, I believe that users should ensure that the input format is correct (i.e., in PyG format), and we do not need to consider the possibility of diverse inputs.

Otherwise, we would have also to do more refactoring beyond the `coalesce` function.
rusty1s: Hey @EdisonLeeeee, we have a use-case where we want to use PyG utilities but have data in `(row, col)` format and would like to avoid the (unnecessary) conversion to a contiguous tensor. That‘s why I suggested to support either case here. I am happy to create an issue.

I do agree that it might make our input arguments quite polluted, so for now I would see this as an internal feature that does not necessarily need to be exposed via type hints. Does that work for you?
EdisonLeeeee: Oh, didn't know that before. Thank you for the clarification. That makes sense :)

Patch:
diff --git a/torch_geometric/utils/coalesce.py b/torch_geometric/utils/coalesce.py
index 20e1b899d..d5305d5aa 100644
--- a/torch_geometric/utils/coalesce.py
+++ b/torch_geometric/utils/coalesce.py
@@ -41,9 +41,9 @@ def coalesce(
     together according to the given :obj:`reduce` option.
 
     Args:
-        edge_index (LongTensor): The edge indices.
-        edge_attr (Tensor or List[Tensor], optional): Edge weights or multi-
-            dimensional edge features.
+        edge_index (torch.Tensor): The edge indices.
+        edge_attr (torch.Tensor or List[torch.Tensor], optional): Edge weights
+            or multi-dimensional edge features.
             If given as a list, will re-shuffle and remove duplicates for all
             its entries. (default: :obj:`None`)
         num_nodes (int, optional): The number of nodes, *i.e.*
@@ -90,17 +90,22 @@ def coalesce(
                 [3, 1, 2]]),
         tensor([1., 1., 1.]))
     """
-    nnz = edge_index.size(1)
+    num_edges = edge_index[0].size(0)
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
-    idx = edge_index.new_empty(nnz + 1)
+    idx = edge_index[0].new_empty(num_edges + 1)
     idx[0] = -1
     idx[1:] = edge_index[1 - int(sort_by_row)]
     idx[1:].mul_(num_nodes).add_(edge_index[int(sort_by_row)])
 
     if not is_sorted:
         idx[1:], perm = index_sort(idx[1:], max_value=num_nodes * num_nodes)
-        edge_index = edge_index[:, perm]
+        if isinstance(edge_index, Tensor):
+            edge_index = edge_index[:, perm]
+        elif isinstance(edge_index, tuple):
+            edge_index = (edge_index[0][perm], edge_index[1][perm])
+        else:
+            raise NotImplementedError
         if isinstance(edge_attr, Tensor):
             edge_attr = edge_attr[perm]
         elif isinstance(edge_attr, (list, tuple)):
@@ -114,12 +119,17 @@ def coalesce(
             return edge_index, edge_attr
         return edge_index
 
-    edge_index = edge_index[:, mask]
+    if isinstance(edge_index, Tensor):
+        edge_index = edge_index[:, mask]
+    elif isinstance(edge_index, tuple):
+        edge_index = (edge_index[0][mask], edge_index[1][mask])
+    else:
+        raise NotImplementedError
 
     dim_size: Optional[int] = None
     if isinstance(edge_attr, (Tensor, list, tuple)) and len(edge_attr) > 0:
         dim_size = edge_index.size(1)
-        idx = torch.arange(0, nnz, device=edge_index.device)
+        idx = torch.arange(0, num_edges, device=edge_index.device)
         idx.sub_(mask.logical_not_().cumsum(dim=0))
 
     if edge_attr is None:
diff --git a/torch_geometric/utils/num_nodes.py b/torch_geometric/utils/num_nodes.py
index c8c4f6388..112f3bb81 100644
--- a/torch_geometric/utils/num_nodes.py
+++ b/torch_geometric/utils/num_nodes.py
@@ -1,5 +1,5 @@
 from copy import copy
-from typing import Dict, Optional, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -14,6 +14,12 @@ def maybe_num_nodes(edge_index, num_nodes):
     pass
 
 
+@torch.jit._overload
+def maybe_num_nodes(edge_index, num_nodes):
+    # type: (Tuple[Tensor, Tensor], Optional[int]) -> int
+    pass
+
+
 @torch.jit._overload
 def maybe_num_nodes(edge_index, num_nodes):
     # type: (SparseTensor, Optional[int]) -> int
@@ -21,7 +27,7 @@ def maybe_num_nodes(edge_index, num_nodes):
 
 
 def maybe_num_nodes(
-    edge_index: Union[Tensor, SparseTensor],
+    edge_index: Union[Tensor, Tuple[Tensor, Tensor], SparseTensor],
     num_nodes: Optional[int] = None,
 ) -> int:
     if num_nodes is not None:
@@ -39,8 +45,15 @@ def maybe_num_nodes(
             return tmp.max() + 1
 
         return int(edge_index.max()) + 1 if edge_index.numel() > 0 else 0
-    else:
+    elif isinstance(edge_index, tuple):
+        return max(
+            int(edge_index[0].max()) + 1 if edge_index[0].numel() > 0 else 0,
+            int(edge_index[1].max()) + 1 if edge_index[1].numel() > 0 else 0,
+        )
+    elif isinstance(edge_index, SparseTensor):
         return max(edge_index.size(0), edge_index.size(1))
+    else:
+        raise NotImplementedError
 
 
 def maybe_num_nodes_dict(
diff --git a/torch_geometric/utils/sort_edge_index.py b/torch_geometric/utils/sort_edge_index.py
index d09d0e230..7be4544a5 100644
--- a/torch_geometric/utils/sort_edge_index.py
+++ b/torch_geometric/utils/sort_edge_index.py
@@ -37,9 +37,9 @@ def sort_edge_index(  # noqa
     """Row-wise sorts :obj:`edge_index`.
 
     Args:
-        edge_index (LongTensor): The edge indices.
-        edge_attr (Tensor or List[Tensor], optional): Edge weights or multi-
-            dimensional edge features.
+        edge_index (torch.Tensor): The edge indices.
+        edge_attr (torch.Tensor or List[torch.Tensor], optional): Edge weights
+            or multi-dimensional edge features.
             If given as a list, will re-shuffle and remove duplicates for all
             its entries. (default: :obj:`None`)
         num_nodes (int, optional): The number of nodes, *i.e.*
@@ -81,7 +81,12 @@ def sort_edge_index(  # noqa
 
     _, perm = index_sort(idx, max_value=num_nodes * num_nodes)
 
-    edge_index = edge_index[:, perm]
+    if isinstance(edge_index, Tensor):
+        edge_index = edge_index[:, perm]
+    elif isinstance(edge_index, tuple):
+        edge_index = (edge_index[0][perm], edge_index[1][perm])
+    else:
+        raise NotImplementedError
 
     if edge_attr is None:
         return edge_index, None


Test patch:
diff --git a/test/utils/test_coalesce.py b/test/utils/test_coalesce.py
index 6d0f9f69b..5c3885df3 100644
--- a/test/utils/test_coalesce.py
+++ b/test/utils/test_coalesce.py
@@ -26,6 +26,11 @@ def test_coalesce():
     assert out[1][0].tolist() == [[4], [3], [2], [6]]
     assert out[1][1].tolist() == [4, 3, 2, 6]
 
+    out = coalesce((edge_index[0], edge_index[1]))
+    assert isinstance(out, tuple)
+    assert out[0].tolist() == [0, 1, 1, 2]
+    assert out[1].tolist() == [1, 0, 2, 1]
+
 
 def test_coalesce_without_duplicates():
     edge_index = torch.tensor([[2, 1, 1, 0], [1, 2, 0, 1]])
diff --git a/test/utils/test_sort_edge_index.py b/test/utils/test_sort_edge_index.py
index d1f5e47fa..b1f11b9cb 100644
--- a/test/utils/test_sort_edge_index.py
+++ b/test/utils/test_sort_edge_index.py
@@ -13,6 +13,11 @@ 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[0], edge_index[1]))
+    assert isinstance(out, tuple)
+    assert out[0].tolist() == [0, 1, 1, 2]
+    assert out[1].tolist() == [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