Repository: pyg-team/pytorch_geometric
Pull request #8641 (2023-12-19 18:05:13)

Expose `EdgeIndex` in `Planetoid`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8641?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`d41de53`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/d41de53b8d3e273acb6fd3d7747139b5534ae022?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.19% compared to head [(`df5aafb`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8641?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.21%.
> Report is 1 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #8641      +/-   ##
==========================================
+ Coverage   89.19%   89.21%   +0.01%     
==========================================
  Files         481      481              
  Lines       30709    30740      +31     
==========================================
+ Hits        27392    27424      +32     
+ Misses       3317     3316       -1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8641?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 139aff97b..fd5d3d093 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -960,17 +960,20 @@ def contiguous(tensor: EdgeIndex) -> EdgeIndex:
 def cat(
     tensors: List[Union[EdgeIndex, Tensor]],
     dim: int = 0,
+    *,
+    out: Optional[Tensor] = None,
 ) -> Union[EdgeIndex, Tensor]:
 
     if len(tensors) == 1:
         return tensors[0]
 
-    out = Tensor.__torch_function__(torch.cat, (Tensor, ), (tensors, dim))
+    output = Tensor.__torch_function__(torch.cat, (Tensor, ), (tensors, dim),
+                                       dict(out=out))
 
     if dim != 1 and dim != -1:  # No valid `EdgeIndex` anymore.
-        return out
+        return output
 
-    out = out.as_subclass(EdgeIndex)
+    output = output.as_subclass(EdgeIndex)
 
     # Post-process `sparse_size`:
     num_rows: Optional[int] = 0
@@ -989,7 +992,7 @@ def cat(
         assert isinstance(num_cols, int)
         num_cols = max(num_cols, tensor.num_cols)
 
-    out._sparse_size = (num_rows, num_cols)
+    output._sparse_size = (num_rows, num_cols)
 
     # Post-process `is_undirected`:
     is_undirected = True
@@ -999,9 +1002,9 @@ def cat(
         else:
             is_undirected = False
 
-    out._is_undirected = is_undirected
+    output._is_undirected = is_undirected
 
-    return out
+    return output
 
 
 @implements(torch.flip)
diff --git a/torch_geometric/io/planetoid.py b/torch_geometric/io/planetoid.py
index 863b2fb3d..a929c7b4f 100644
--- a/torch_geometric/io/planetoid.py
+++ b/torch_geometric/io/planetoid.py
@@ -1,12 +1,13 @@
 import os.path as osp
 import warnings
 from itertools import repeat
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import fsspec
 import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data import Data
 from torch_geometric.io import read_txt_array
 from torch_geometric.typing import SparseTensor
@@ -112,17 +113,25 @@ def read_file(folder: str, prefix: str, name: str) -> Tensor:
 def edge_index_from_dict(
     graph_dict: Dict[int, List[int]],
     num_nodes: Optional[int] = None,
-) -> Tensor:
+) -> EdgeIndex:
     rows: List[int] = []
     cols: List[int] = []
     for key, value in graph_dict.items():
         rows += repeat(key, len(value))
         cols += value
-    edge_index = torch.stack([torch.tensor(rows), torch.tensor(cols)], dim=0)
+    row = torch.tensor(rows)
+    col = torch.tensor(cols)
+
+    edge_index: Union[EdgeIndex, Tensor] = EdgeIndex(
+        torch.stack([row, col], dim=0),
+        is_undirected=True,
+        sparse_size=(num_nodes, num_nodes),
+    )
 
     # NOTE: There are some duplicated edges and self loops in the datasets.
     #       Other implementations do not remove them!
     edge_index, _ = remove_self_loops(edge_index)
-    edge_index = coalesce(edge_index, num_nodes=num_nodes)
+    edge_index = coalesce(edge_index, num_nodes=num_nodes, sort_by_row=False)
+    assert isinstance(edge_index, EdgeIndex)
 
     return edge_index
diff --git a/torch_geometric/utils/_coalesce.py b/torch_geometric/utils/_coalesce.py
index 5b026b872..6ffb3faac 100644
--- a/torch_geometric/utils/_coalesce.py
+++ b/torch_geometric/utils/_coalesce.py
@@ -4,6 +4,8 @@ from typing import List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
+from torch_geometric.edge_index import SortOrder
 from torch_geometric.typing import OptTensor
 from torch_geometric.utils import index_sort, scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
@@ -133,6 +135,10 @@ def coalesce(  # noqa: F811
     idx[1:] = edge_index[1 - int(sort_by_row)]
     idx[1:].mul_(num_nodes).add_(edge_index[int(sort_by_row)])
 
+    is_undirected = False
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        is_undirected = edge_index.is_undirected
+
     if not is_sorted:
         idx[1:], perm = index_sort(idx[1:], max_value=num_nodes * num_nodes)
         if isinstance(edge_index, Tensor):
@@ -146,6 +152,10 @@ def coalesce(  # noqa: F811
         elif isinstance(edge_attr, (list, tuple)):
             edge_attr = [e[perm] for e in edge_attr]
 
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        edge_index._sort_order = SortOrder('row' if sort_by_row else 'col')
+        edge_index._is_undirected = is_undirected
+
     mask = idx[1:] > idx[:-1]
 
     # Only perform expensive merging in case there exists duplicates:
@@ -158,6 +168,8 @@ def coalesce(  # noqa: F811
 
     if isinstance(edge_index, Tensor):
         edge_index = edge_index[:, mask]
+        if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+            edge_index._is_undirected = is_undirected
     elif isinstance(edge_index, tuple):
         edge_index = (edge_index[0][mask], edge_index[1][mask])
     else:
diff --git a/torch_geometric/utils/_sort_edge_index.py b/torch_geometric/utils/_sort_edge_index.py
index b7b233197..e36686dbc 100644
--- a/torch_geometric/utils/_sort_edge_index.py
+++ b/torch_geometric/utils/_sort_edge_index.py
@@ -1,8 +1,11 @@
 import typing
 from typing import List, Optional, Tuple, Union
 
+import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
+from torch_geometric.edge_index import SortOrder
 from torch_geometric.typing import OptTensor
 from torch_geometric.utils import index_sort
 from torch_geometric.utils.num_nodes import maybe_num_nodes
@@ -108,7 +111,13 @@ def sort_edge_index(  # noqa: F811
     _, perm = index_sort(idx, max_value=num_nodes * num_nodes)
 
     if isinstance(edge_index, Tensor):
+        is_undirected = False
+        if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+            is_undirected = edge_index.is_undirected
         edge_index = edge_index[:, perm]
+        if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+            edge_index._sort_order = SortOrder('row' if sort_by_row else 'col')
+            edge_index._is_undirected = is_undirected
     elif isinstance(edge_index, tuple):
         edge_index = (edge_index[0][perm], edge_index[1][perm])
     else:
diff --git a/torch_geometric/utils/loop.py b/torch_geometric/utils/loop.py
index a5ac150b3..c76c78c40 100644
--- a/torch_geometric/utils/loop.py
+++ b/torch_geometric/utils/loop.py
@@ -4,6 +4,7 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
 from torch_geometric.utils import scatter
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 from torch_geometric.utils.sparse import (
@@ -104,9 +105,16 @@ def remove_self_loops(  # noqa: F811
         size = (edge_index.size(0), edge_index.size(1))
         edge_index, value = to_edge_index(edge_index)
 
+    is_undirected = False
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        is_undirected = edge_index.is_undirected
+
     mask = edge_index[0] != edge_index[1]
     edge_index = edge_index[:, mask]
 
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        edge_index._is_undirected = is_undirected
+
     if layout is not None:
         assert edge_attr is None
         assert value is not None


Test patch:
diff --git a/test/datasets/test_planetoid.py b/test/datasets/test_planetoid.py
index 05dd9ebbc..31bef8d7c 100644
--- a/test/datasets/test_planetoid.py
+++ b/test/datasets/test_planetoid.py
@@ -1,3 +1,4 @@
+from torch_geometric import EdgeIndex
 from torch_geometric.loader import DataLoader
 from torch_geometric.testing import onlyOnline
 
@@ -50,6 +51,11 @@ def test_citeseer_with_random_split(get_dataset):
         num_test=41,
     )
     data = dataset[0]
+    assert isinstance(data.edge_index, EdgeIndex)
+    assert data.edge_index.sparse_size() == (data.num_nodes, data.num_nodes)
+    assert data.edge_index.is_undirected
+    assert data.edge_index.is_sorted_by_col
+
     assert data.train_mask.sum() == dataset.num_classes * 11
     assert data.val_mask.sum() == 29
     assert data.test_mask.sum() == 41
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 484a524b9..a7b94da0c 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -381,6 +381,12 @@ def test_cat(dtype, device, is_undirected):
     assert out.size() == (4, 4)
     assert not isinstance(out, EdgeIndex)
 
+    inplace = torch.empty(2, 8, dtype=dtype, device=device)
+    out = torch.cat([adj1, adj2], dim=1, out=inplace)
+    assert out.data_ptr() == inplace.data_ptr()
+    assert isinstance(out, EdgeIndex)
+    assert not isinstance(inplace, EdgeIndex)
+
 
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)