Repository: pyg-team/pytorch_geometric
Pull request #8686 (2023-12-29 18:30:02)

`EdgeIndex` integration in `NeighborLoader`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`a218de3`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/a218de306f2027ed8aabe35306def4667dc54401?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.48% compared to head [(`314d845`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.48%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2VkZ2VfaW5kZXgucHk=) | 88.88% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

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


```diff
@@            Coverage Diff             @@
##           master    #8686      +/-   ##
==========================================
- Coverage   89.48%   89.48%   -0.01%     
==========================================
  Files         481      481              
  Lines       31009    31027      +18     
==========================================
+ Hits        27748    27764      +16     
- Misses       3261     3263       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8686?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 4d3b04dbc..46bd3e771 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -934,6 +934,45 @@ class SortReturnType(NamedTuple):
     indices: Union[Tensor, slice]
 
 
+@implements(Tensor.__repr__)
+def __repr__(
+    tensor: EdgeIndex,
+    *,
+    tensor_contents: Optional[str] = None,
+) -> str:
+    # Monkey-patch `torch._tensor_str._add_suffixes`. There might exist better
+    # solutions to attach additional metadata, but this seems to be the most
+    # straightforward one to inherit most of the `torch.Tensor` print logic:
+    orig_fn = torch._tensor_str._add_suffixes
+
+    def _add_suffixes(
+        tensor_str: str,
+        suffixes: List[str],
+        indent: int,
+        force_newline: bool,
+    ) -> str:
+
+        num_rows, num_cols = tensor.sparse_size()
+        if num_rows is not None or num_cols is not None:
+            size_repr = f"({num_rows or '?'}, {num_cols or '?'})"
+            suffixes.append(f'sparse_size={size_repr}')
+
+        suffixes.append(f'nnz={tensor.size(1)}')
+
+        if tensor.is_sorted:
+            suffixes.append(f'sort_order={tensor.sort_order}')
+
+        if tensor.is_undirected:
+            suffixes.append('is_undirected=True')
+
+        return orig_fn(tensor_str, suffixes, indent, force_newline)
+
+    torch._tensor_str._add_suffixes = _add_suffixes
+    out = torch._tensor_str._str(tensor, tensor_contents=tensor_contents)
+    torch._tensor_str._add_suffixes = orig_fn
+    return out
+
+
 def apply_(
     tensor: EdgeIndex,
     fn: Callable,
diff --git a/torch_geometric/loader/utils.py b/torch_geometric/loader/utils.py
index dac7cde40..70828d0c4 100644
--- a/torch_geometric/loader/utils.py
+++ b/torch_geometric/loader/utils.py
@@ -8,6 +8,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.data import (
     Data,
     FeatureStore,
@@ -104,8 +105,13 @@ def filter_edge_store_(store: EdgeStorage, out_store: EdgeStorage, row: Tensor,
     # which represents the new graph as denoted by `(row, col)`:
     for key, value in store.items():
         if key == 'edge_index':
-            edge_index = torch.stack([row, col], dim=0)
-            out_store.edge_index = edge_index.to(value.device)
+            # TODO Integrate `EdgeIndex` into `custom_store`.
+            out_store.edge_index = EdgeIndex(
+                torch.stack([row, col], dim=0).to(value.device),
+                sparse_size=out_store.size(),
+                sort_order='col',
+                # TODO Support `is_undirected`.
+            )
 
         elif key == 'adj_t':
             # NOTE: We expect `(row, col)` to be sorted by `col` (CSC layout).
diff --git a/torch_geometric/sampler/utils.py b/torch_geometric/sampler/utils.py
index 402d794e6..584eec187 100644
--- a/torch_geometric/sampler/utils.py
+++ b/torch_geometric/sampler/utils.py
@@ -135,7 +135,12 @@ def to_bidirectional(
     if edge_id is not None:
         edge_id = torch.cat([edge_id, rev_edge_id], dim=0)
 
-    (row, col), edge_id = coalesce(edge_index, edge_id, reduce='any')
+    (row, col), edge_id = coalesce(
+        edge_index,
+        edge_id,
+        sort_by_row=False,
+        reduce='any',
+    )
 
     return row, col, edge_id
 


Test patch:
diff --git a/test/loader/test_neighbor_loader.py b/test/loader/test_neighbor_loader.py
index 04bba6b4e..632e92a77 100644
--- a/test/loader/test_neighbor_loader.py
+++ b/test/loader/test_neighbor_loader.py
@@ -4,6 +4,8 @@ import numpy as np
 import pytest
 import torch
 
+import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import GraphConv, to_hetero
@@ -33,6 +35,23 @@ from torch_geometric.utils import (
     to_undirected,
 )
 
+DTYPES = [
+    pytest.param(torch.int64, id='int64'),
+    pytest.param(torch.int32, id='int32'),
+]
+
+SUBGRAPH_TYPES = [
+    pytest.param(SubgraphType.directional, id='directional'),
+    pytest.param(SubgraphType.bidirectional, id='bidirectional'),
+    pytest.param(SubgraphType.induced, id='induced'),
+]
+
+FILTER_PER_WORKERS = [
+    pytest.param(None, id='auto_filter'),
+    pytest.param(True, id='filter_per_worker'),
+    pytest.param(False, id='filter_in_main'),
+]
+
 
 def is_subset(subedge_index, edge_index, src_idx, dst_idx):
     num_nodes = int(edge_index.max()) + 1
@@ -44,15 +63,21 @@ def is_subset(subedge_index, edge_index, src_idx, dst_idx):
 
 @withCUDA
 @onlyNeighborSampler
-@pytest.mark.parametrize('subgraph_type', list(SubgraphType))
-@pytest.mark.parametrize('dtype', [torch.int64, torch.int32])
-@pytest.mark.parametrize('filter_per_worker', [None, True, False])
-def test_homo_neighbor_loader_basic(device, subgraph_type, dtype,
-                                    filter_per_worker):
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES)
+@pytest.mark.parametrize('filter_per_worker', FILTER_PER_WORKERS)
+def test_homo_neighbor_loader_basic(
+    device,
+    subgraph_type,
+    dtype,
+    filter_per_worker,
+):
+    if dtype != torch.int64 and not torch_geometric.typing.WITH_PT20:
+        return
+    induced = SubgraphType.induced
     if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE:
         return
-    if (dtype != torch.int64
-            and (not WITH_PYG_LIB or subgraph_type == SubgraphType.induced)):
+    if dtype != torch.int64 and (not WITH_PYG_LIB or subgraph_type == induced):
         return
 
     torch.manual_seed(12345)
@@ -85,6 +110,11 @@ def test_homo_neighbor_loader_basic(device, subgraph_type, dtype,
         assert batch.n_id.size() == (batch.num_nodes, )
         assert batch.input_id.numel() == batch.batch_size == 20
         assert batch.x.min() >= 0 and batch.x.max() < 100
+        assert isinstance(batch.edge_index, EdgeIndex)
+        batch.edge_index.validate()
+        size = (batch.num_nodes, batch.num_nodes)
+        assert batch.edge_index.sparse_size() == size
+        assert batch.edge_index.sort_order == 'col'
         assert batch.edge_index.device == device
         assert batch.edge_index.min() >= 0
         assert batch.edge_index.max() < batch.num_nodes
@@ -112,13 +142,15 @@ def test_homo_neighbor_loader_basic(device, subgraph_type, dtype,
 
 
 @onlyNeighborSampler
-@pytest.mark.parametrize('subgraph_type', list(SubgraphType))
-@pytest.mark.parametrize('dtype', [torch.int64, torch.int32])
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES)
 def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
+    if dtype != torch.int64 and not torch_geometric.typing.WITH_PT20:
+        return
+    induced = SubgraphType.induced
     if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE:
         return
-    if (dtype != torch.int64
-            and (not WITH_PYG_LIB or subgraph_type == SubgraphType.induced)):
+    if dtype != torch.int64 and (not WITH_PYG_LIB or subgraph_type == induced):
         return
 
     torch.manual_seed(12345)
@@ -190,6 +222,14 @@ def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
                                          ('paper', 'to', 'author'),
                                          ('author', 'to', 'paper')}
 
+        for edge_type, edge_index in batch.edge_index_dict.items():
+            src, _, dst = edge_type
+            assert isinstance(edge_index, EdgeIndex)
+            edge_index.validate()
+            size = (batch[src].num_nodes, batch[dst].num_nodes)
+            assert edge_index.sparse_size() == size
+            assert edge_index.sort_order == 'col'
+
         row, col = batch['paper', 'paper'].edge_index
         assert row.min() >= 0 and row.max() < batch['paper'].num_nodes
         assert col.min() >= 0 and col.max() < batch['paper'].num_nodes
@@ -230,9 +270,11 @@ def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
             assert 'e_id' not in batch['paper', 'author']
             assert 'edge_attr' not in batch['paper', 'author']
 
+            edge_index1 = batch['paper', 'author'].edge_index
+            edge_index2 = batch['author', 'paper'].edge_index
             assert torch.equal(
-                batch['paper', 'author'].edge_index,
-                sort_edge_index(batch['author', 'paper'].edge_index.flip([0])),
+                edge_index1,
+                sort_edge_index(edge_index2.flip([0]), sort_by_row=False),
             )
 
         row, col = batch['author', 'paper'].edge_index
@@ -254,9 +296,11 @@ def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
             assert 'e_id' not in batch['author', 'paper']
             assert 'edge_attr' not in batch['author', 'paper']
 
+            edge_index1 = batch['author', 'paper'].edge_index
+            edge_index2 = batch['paper', 'author'].edge_index
             assert torch.equal(
-                batch['author', 'paper'].edge_index,
-                sort_edge_index(batch['paper', 'author'].edge_index.flip([0])),
+                edge_index1,
+                sort_edge_index(edge_index2.flip([0]), sort_by_row=False),
             )
 
         # Test for isolated nodes (there shouldn't exist any):
@@ -265,7 +309,7 @@ def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
 
 @onlyOnline
 @onlyNeighborSampler
-@pytest.mark.parametrize('subgraph_type', list(SubgraphType))
+@pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES)
 def test_homo_neighbor_loader_on_cora(get_dataset, subgraph_type):
     if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE:
         return
@@ -311,7 +355,7 @@ def test_homo_neighbor_loader_on_cora(get_dataset, subgraph_type):
 
 @onlyOnline
 @onlyNeighborSampler
-@pytest.mark.parametrize('subgraph_type', list(SubgraphType))
+@pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES)
 def test_hetero_neighbor_loader_on_cora(get_dataset, subgraph_type):
     if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE:
         return
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index f8b567320..7d7083df0 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -47,10 +47,17 @@ def test_basic(dtype, device):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], **kwargs)
     adj.validate()
     assert isinstance(adj, EdgeIndex)
+
+    dtype_str = '' if dtype == torch.long else f', dtype={dtype}'
     if torch_geometric.typing.WITH_PT112:
-        assert str(adj).startswith('EdgeIndex([[0, 1, 1, 2],')
+        assert str(adj) == (
+            f'EdgeIndex([[0, 1, 1, 2],\n'
+            f'           [1, 0, 2, 1]]{dtype_str}, sparse_size=(3, 3), nnz=4)')
     else:
-        assert str(adj).startswith('tensor([[0, 1, 1, 2],')
+        assert str(adj) == (
+            f'tensor([[0, 1, 1, 2],\n'
+            f'        [1, 0, 2, 1]]{dtype_str}, sparse_size=(3, 3), nnz=4)')
+
     assert adj.dtype == dtype
     assert adj.device == device
     assert adj.sparse_size() == (3, 3)