Repository: pyg-team/pytorch_geometric
Pull request #6879 (2023-03-08 12:09:29)

Changed the interface of `coalesce` 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/6879?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6879](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5d19333) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/581dd016c568f33b1a38ac4fe2c79ac27b10d8c5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (581dd01) will **increase** coverage by `0.00%`.
> The diff coverage is `95.23%`.

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

```diff
@@           Coverage Diff           @@
##           master    #6879   +/-   ##
=======================================
  Coverage   91.53%   91.53%           
=======================================
  Files         427      427           
  Lines       23173    23166    -7     
=======================================
- Hits        21211    21205    -6     
+ Misses       1962     1961    -1     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?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/undirected.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3VuZGlyZWN0ZWQucHk=) | `90.24% <66.66%> (-0.67%)` | :arrow_down: |
| [torch\_geometric/io/tu.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2lvL3R1LnB5) | `90.58% <100.00%> (-0.22%)` | :arrow_down: |
| [torch\_geometric/transforms/line\_graph.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvbGluZV9ncmFwaC5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/transforms/two\_hop.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvdHdvX2hvcC5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/coalesce.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?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%> (+2.32%)` | :arrow_up: |
| [torch\_geometric/utils/grid.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2dyaWQucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/sort\_edge\_index.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6879?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 31389e226..665ee53fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -98,7 +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))
+- Changed the interface of `sort_edge_index`, `coalesce` and `to_undirected` 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), [#6879](https://github.com/pyg-team/pytorch_geometric/issues/6879))
 - 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/datasets/snap_dataset.py b/torch_geometric/datasets/snap_dataset.py
index 44637ae19..12d54c880 100644
--- a/torch_geometric/datasets/snap_dataset.py
+++ b/torch_geometric/datasets/snap_dataset.py
@@ -129,7 +129,7 @@ def read_soc(files: List[str], name: str) -> List[Data]:
                              skiprows=skiprows, dtype=np.int64)
     edge_index = torch.from_numpy(edge_index.values).t()
     num_nodes = edge_index.max().item() + 1
-    edge_index, _ = coalesce(edge_index, None, num_nodes, num_nodes)
+    edge_index = coalesce(edge_index, num_nodes=num_nodes)
 
     return [Data(edge_index=edge_index, num_nodes=num_nodes)]
 
@@ -147,7 +147,7 @@ def read_wiki(files: List[str], name: str) -> List[Data]:
 
     edge_index = idx_assoc[edge_index]
     num_nodes = edge_index.max().item() + 1
-    edge_index, _ = coalesce(edge_index, None, num_nodes, num_nodes)
+    edge_index = coalesce(edge_index, num_nodes=num_nodes)
 
     return [Data(edge_index=edge_index, num_nodes=num_nodes)]
 
diff --git a/torch_geometric/io/tu.py b/torch_geometric/io/tu.py
index 359f049f6..01b36025e 100644
--- a/torch_geometric/io/tu.py
+++ b/torch_geometric/io/tu.py
@@ -67,10 +67,7 @@ def read_tu_data(folder, prefix):
 
     num_nodes = edge_index.max().item() + 1 if x is None else x.size(0)
     edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
-    if edge_attr is None:
-        edge_index = coalesce(edge_index, num_nodes=num_nodes)
-    else:
-        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes)
+    edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes)
 
     data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
     data, slices = split(data, batch)
diff --git a/torch_geometric/transforms/line_graph.py b/torch_geometric/transforms/line_graph.py
index 9d63ae388..531e13ce3 100644
--- a/torch_geometric/transforms/line_graph.py
+++ b/torch_geometric/transforms/line_graph.py
@@ -37,10 +37,7 @@ class LineGraph(BaseTransform):
     def __call__(self, data: Data) -> Data:
         N = data.num_nodes
         edge_index, edge_attr = data.edge_index, data.edge_attr
-        if edge_attr is None:
-            edge_index = coalesce(edge_index, num_nodes=N)
-        else:
-            edge_index, edge_attr = coalesce(edge_index, edge_attr, N)
+        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes=N)
         row, col = edge_index
 
         if self.force_directed or data.is_directed():
diff --git a/torch_geometric/transforms/two_hop.py b/torch_geometric/transforms/two_hop.py
index 16ff5abb8..a05792ec1 100644
--- a/torch_geometric/transforms/two_hop.py
+++ b/torch_geometric/transforms/two_hop.py
@@ -23,14 +23,13 @@ class TwoHop(BaseTransform):
         edge_index2, _ = remove_self_loops(edge_index2)
 
         edge_index = torch.cat([edge_index, edge_index2], dim=1)
-        if edge_attr is None:
-            data.edge_index = coalesce(edge_index, num_nodes=N)
-        else:
+
+        if edge_attr is not None:
             # We treat newly added edge features as "zero-features":
             edge_attr2 = edge_attr.new_zeros(edge_index2.size(1),
                                              *edge_attr.size()[1:])
             edge_attr = torch.cat([edge_attr, edge_attr2], dim=0)
-            edge_index, edge_attr = coalesce(edge_index, edge_attr, N)
-            data.edge_index, data.edge_attr = edge_index, edge_attr
+
+        data.edge_index, data.edge_attr = coalesce(edge_index, edge_attr, N)
 
         return data
diff --git a/torch_geometric/utils/coalesce.py b/torch_geometric/utils/coalesce.py
index f85a05e3b..66b4c845d 100644
--- a/torch_geometric/utils/coalesce.py
+++ b/torch_geometric/utils/coalesce.py
@@ -3,40 +3,40 @@ from typing import List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import OptTensor
 from torch_geometric.utils import index_sort, scatter
 
 from .num_nodes import maybe_num_nodes
 
+MISSING = '???'
+
 
 @torch.jit._overload
-def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
-             is_sorted=False, sort_by_row=True):
-    # type: (Tensor, Optional[bool], Optional[int], str, bool, bool) -> Tensor  # noqa
+def coalesce(edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
+    # type: (Tensor, str, Optional[int], str, bool, bool) -> Tensor  # noqa
     pass
 
 
 @torch.jit._overload
-def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
-             is_sorted=False, sort_by_row=True):
-    # type: (Tensor, Tensor, Optional[int], str, bool, bool) -> Tuple[Tensor, Tensor]  # noqa
+def coalesce(edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
+    # type: (Tensor, Optional[Tensor], Optional[int], str, bool, bool) -> Tuple[Tensor, Optional[Tensor]]  # noqa
     pass
 
 
 @torch.jit._overload
-def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
-             is_sorted=False, sort_by_row=True):
+def coalesce(edge_index, edge_attr, num_nodes, reduce, is_sorted, sort_by_row):
     # type: (Tensor, List[Tensor], Optional[int], str, bool, bool) -> Tuple[Tensor, List[Tensor]]  # noqa
     pass
 
 
 def coalesce(
     edge_index: Tensor,
-    edge_attr: Union[Optional[Tensor], List[Tensor]] = None,
+    edge_attr: Union[OptTensor, List[Tensor], str] = MISSING,
     num_nodes: Optional[int] = None,
-    reduce: str = "add",
+    reduce: str = 'add',
     is_sorted: bool = False,
     sort_by_row: bool = True,
-) -> Union[Tensor, Tuple[Tensor, Tensor], Tuple[Tensor, List[Tensor]]]:
+) -> Union[Tensor, Tuple[Tensor, OptTensor], Tuple[Tensor, List[Tensor]]]:
     """Row-wise sorts :obj:`edge_index` and removes its duplicated entries.
     Duplicate entries in :obj:`edge_attr` are merged by scattering them
     together according to the given :obj:`reduce` option.
@@ -57,8 +57,8 @@ def coalesce(
         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]]`)
 
     Example:
 
@@ -105,23 +105,24 @@ def coalesce(
 
     # Only perform expensive merging in case there exists duplicates:
     if mask.all():
-        if isinstance(edge_attr, (Tensor, list, tuple)):
+        if edge_attr is None or isinstance(edge_attr, (Tensor, list, tuple)):
             return edge_index, edge_attr
         return edge_index
 
     edge_index = edge_index[:, mask]
 
-    if edge_attr is None:
-        return edge_index
-
-    dim_size = edge_index.size(1)
-    idx = torch.arange(0, nnz, device=edge_index.device)
-    idx.sub_(mask.logical_not_().cumsum(dim=0))
+    dim_size: Optional[int] = None
+    if isinstance(edge_attr, (Tensor, list, tuple)):
+        dim_size = edge_index.size(1)
+        idx = torch.arange(0, nnz, device=edge_index.device)
+        idx.sub_(mask.logical_not_().cumsum(dim=0))
 
+    if edge_attr is None:
+        return edge_index, None
     if isinstance(edge_attr, Tensor):
         edge_attr = scatter(edge_attr, idx, 0, dim_size, reduce)
         return edge_index, edge_attr
-    elif isinstance(edge_attr, (list, tuple)):
+    if isinstance(edge_attr, (list, tuple)):
         edge_attr = [scatter(e, idx, 0, dim_size, reduce) for e in edge_attr]
         return edge_index, edge_attr
 
diff --git a/torch_geometric/utils/grid.py b/torch_geometric/utils/grid.py
index 9393c911b..c94cc1c53 100644
--- a/torch_geometric/utils/grid.py
+++ b/torch_geometric/utils/grid.py
@@ -64,7 +64,7 @@ def grid_index(
     row, col = row[mask], col[mask]
 
     edge_index = torch.stack([row, col], dim=0)
-    edge_index = coalesce(edge_index, None, height * width)
+    edge_index = coalesce(edge_index, num_nodes=height * width)
     return edge_index
 
 
diff --git a/torch_geometric/utils/sort_edge_index.py b/torch_geometric/utils/sort_edge_index.py
index 63b2f2b36..ab83a341e 100644
--- a/torch_geometric/utils/sort_edge_index.py
+++ b/torch_geometric/utils/sort_edge_index.py
@@ -3,6 +3,7 @@ from typing import List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import OptTensor
 from torch_geometric.utils import index_sort
 
 from .num_nodes import maybe_num_nodes
@@ -30,10 +31,10 @@ def sort_edge_index(edge_index, edge_attr, num_nodes, sort_by_row):
 
 def sort_edge_index(
     edge_index: Tensor,
-    edge_attr: Union[Optional[Tensor], List[Tensor], str] = MISSING,
+    edge_attr: Union[OptTensor, List[Tensor], str] = MISSING,
     num_nodes: Optional[int] = None,
     sort_by_row: bool = True,
-) -> Union[Tensor, Tuple[Tensor, Tensor], Tuple[Tensor, List[Tensor]]]:
+) -> Union[Tensor, Tuple[Tensor, OptTensor], Tuple[Tensor, List[Tensor]]]:
     """Row-wise sorts :obj:`edge_index`.
 
     Args:
@@ -80,7 +81,7 @@ def sort_edge_index(
         return edge_index, None
     if isinstance(edge_attr, Tensor):
         return edge_index, edge_attr[perm]
-    elif isinstance(edge_attr, (list, tuple)):
+    if isinstance(edge_attr, (list, tuple)):
         return edge_index, [e[perm] for e in edge_attr]
 
     return edge_index
diff --git a/torch_geometric/utils/undirected.py b/torch_geometric/utils/undirected.py
index 591477e24..0fd5e3979 100644
--- a/torch_geometric/utils/undirected.py
+++ b/torch_geometric/utils/undirected.py
@@ -3,26 +3,29 @@ from typing import List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import OptTensor
 from torch_geometric.utils import coalesce, sort_edge_index
 
 from .num_nodes import maybe_num_nodes
 
+MISSING = '???'
+
 
 @torch.jit._overload
-def is_undirected(edge_index, edge_attr=None, num_nodes=None):
+def is_undirected(edge_index, edge_attr, num_nodes):
     # type: (Tensor, Optional[Tensor], Optional[int]) -> bool  # noqa
     pass
 
 
 @torch.jit._overload
-def is_undirected(edge_index, edge_attr=None, num_nodes=None):
+def is_undirected(edge_index, edge_attr, num_nodes):
     # type: (Tensor, List[Tensor], Optional[int]) -> bool  # noqa
     pass
 
 
 def is_undirected(
     edge_index: Tensor,
-    edge_attr: Union[Optional[Tensor], List[Tensor]] = None,
+    edge_attr: Union[OptTensor, List[Tensor]] = None,
     num_nodes: Optional[int] = None,
 ) -> bool:
     r"""Returns :obj:`True` if the graph given by :attr:`edge_index` is
@@ -84,29 +87,29 @@ def is_undirected(
 
 
 @torch.jit._overload
-def to_undirected(edge_index, edge_attr=None, num_nodes=None, reduce="add"):
-    # type: (Tensor, Optional[bool], Optional[int], str) -> Tensor  # noqa
+def to_undirected(edge_index, edge_attr, num_nodes, reduce):
+    # type: (Tensor, str, Optional[int], str) -> Tensor  # noqa
     pass
 
 
 @torch.jit._overload
-def to_undirected(edge_index, edge_attr=None, num_nodes=None, reduce="add"):
-    # type: (Tensor, Tensor, Optional[int], str) -> Tuple[Tensor, Tensor]  # noqa
+def to_undirected(edge_index, edge_attr, num_nodes, reduce):
+    # type: (Tensor, Optional[Tensor], Optional[int], str) -> Tuple[Tensor, Optional[Tensor]]  # noqa
     pass
 
 
 @torch.jit._overload
-def to_undirected(edge_index, edge_attr=None, num_nodes=None, reduce="add"):
+def to_undirected(edge_index, edge_attr, num_nodes, reduce):
     # type: (Tensor, List[Tensor], Optional[int], str) -> Tuple[Tensor, List[Tensor]]  # noqa
     pass
 
 
 def to_undirected(
     edge_index: Tensor,
-    edge_attr: Union[Optional[Tensor], List[Tensor]] = None,
+    edge_attr: Union[OptTensor, List[Tensor], str] = MISSING,
     num_nodes: Optional[int] = None,
-    reduce: str = "add",
-) -> Union[Tensor, Tuple[Tensor, Tensor], Tuple[Tensor, List[Tensor]]]:
+    reduce: str = 'add',
+) -> Union[Tensor, Tuple[OptTensor, Tensor], Tuple[Tensor, List[Tensor]]]:
     r"""Converts the graph given by :attr:`edge_index` to an undirected graph
     such that :math:`(j,i) \in \mathcal{E}` for every edge :math:`(i,j) \in
     \mathcal{E}`.
@@ -123,8 +126,8 @@ def to_undirected(
             features (:obj:`"add"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
             :obj:`"mul"`). (default: :obj:`"add"`)
 
-    :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:
 
@@ -150,7 +153,7 @@ def to_undirected(
     """
     # Maintain backward compatibility to `to_undirected(edge_index, num_nodes)`
     if isinstance(edge_attr, int):
-        edge_attr = None
+        edge_attr = MISSING
         num_nodes = edge_attr
 
     row, col = edge_index[0], edge_index[1]


Test patch:
diff --git a/test/transforms/test_radius_graph.py b/test/transforms/test_radius_graph.py
index a9736e557..0fbf9fd96 100644
--- a/test/transforms/test_radius_graph.py
+++ b/test/transforms/test_radius_graph.py
@@ -1,9 +1,9 @@
 import torch
-from torch_sparse import coalesce
 
 from torch_geometric.data import Data
 from torch_geometric.testing import withPackage
 from torch_geometric.transforms import RadiusGraph
+from torch_geometric.utils import coalesce
 
 
 @withPackage('torch_cluster')
@@ -16,6 +16,5 @@ def test_radius_graph():
     data = RadiusGraph(r=1.5)(data)
     assert len(data) == 2
     assert data.pos.tolist() == pos.tolist()
-    edge_index, _ = coalesce(data.edge_index, None, 6, 6)
-    assert edge_index.tolist() == [[0, 0, 1, 1, 1, 2, 3, 3],
-                                   [1, 3, 0, 2, 3, 1, 0, 1]]
+    assert coalesce(data.edge_index).tolist() == [[0, 0, 1, 1, 1, 2, 3, 3],
+                                                  [1, 3, 0, 2, 3, 1, 0, 1]]
diff --git a/test/utils/test_coalesce.py b/test/utils/test_coalesce.py
new file mode 100644
index 000000000..6d0f9f69b
--- /dev/null
+++ b/test/utils/test_coalesce.py
@@ -0,0 +1,88 @@
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import coalesce
+
+
+def test_coalesce():
+    edge_index = torch.tensor([[2, 1, 1, 0, 2], [1, 2, 0, 1, 1]])
+    edge_attr = torch.tensor([[1], [2], [3], [4], [5]])
+
+    out = coalesce(edge_index)
+    assert out.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+
+    out = coalesce(edge_index, None)
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1] is None
+
+    out = coalesce(edge_index, edge_attr)
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1].tolist() == [[4], [3], [2], [6]]
+
+    out = coalesce(edge_index, [edge_attr, edge_attr.view(-1)])
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1][0].tolist() == [[4], [3], [2], [6]]
+    assert out[1][1].tolist() == [4, 3, 2, 6]
+
+
+def test_coalesce_without_duplicates():
+    edge_index = torch.tensor([[2, 1, 1, 0], [1, 2, 0, 1]])
+    edge_attr = torch.tensor([[1], [2], [3], [4]])
+
+    out = coalesce(edge_index)
+    assert out.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+
+    out = coalesce(edge_index, None)
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1] is None
+
+    out = coalesce(edge_index, edge_attr)
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1].tolist() == [[4], [3], [2], [1]]
+
+    out = coalesce(edge_index, [edge_attr, edge_attr.view(-1)])
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1][0].tolist() == [[4], [3], [2], [1]]
+    assert out[1][1].tolist() == [4, 3, 2, 1]
+
+
+def test_coalesce_jit():
+    @torch.jit.script
+    def wrapper1(edge_index: Tensor) -> Tensor:
+        return coalesce(edge_index)
+
+    @torch.jit.script
+    def wrapper2(
+        edge_index: Tensor,
+        edge_attr: Optional[Tensor],
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        return coalesce(edge_index, edge_attr)
+
+    @torch.jit.script
+    def wrapper3(
+        edge_index: Tensor,
+        edge_attr: List[Tensor],
+    ) -> Tuple[Tensor, List[Tensor]]:
+        return coalesce(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.view(-1)])
+    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.view(-1).size()
diff --git a/test/utils/test_sort_edge_index.py b/test/utils/test_sort_edge_index.py
index 1e57fbfc3..d1f5e47fa 100644
--- a/test/utils/test_sort_edge_index.py
+++ b/test/utils/test_sort_edge_index.py
@@ -21,6 +21,11 @@ def test_sort_edge_index():
     assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
     assert out[1].tolist() == [[4], [3], [2], [1]]
 
+    out = sort_edge_index(edge_index, [edge_attr, edge_attr.view(-1)])
+    assert out[0].tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out[1][0].tolist() == [[4], [3], [2], [1]]
+    assert out[1][1].tolist() == [4, 3, 2, 1]
+
 
 def test_sort_edge_index_jit():
     @torch.jit.script
@@ -55,8 +60,8 @@ def test_sort_edge_index_jit():
     assert out[0].size() == edge_index.size()
     assert out[1].size() == edge_attr.size()
 
-    out = wrapper3(edge_index, [edge_attr, edge_attr])
+    out = wrapper3(edge_index, [edge_attr, edge_attr.view(-1)])
     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()
+    assert out[1][1].size() == edge_attr.view(-1).size()