Repository: pyg-team/pytorch_geometric
Pull request #9275 (2024-05-03 21:15:36)

Move common utils out of `EdgeIndex`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `96.96970%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 87.87%. Comparing base [(`01e8303`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/01e830375c429881c3f717677dbf97dd4d53f2e8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`8d748d2`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).
> Report is 13 commits behind head on master.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?dropdown=coverage&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/index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?src=pr&el=tree&filepath=torch_geometric%2Findex.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2luZGV4LnB5) | 90.00% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?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    #9275      +/-   ##
==========================================
+ Coverage   85.22%   87.87%   +2.64%     
==========================================
  Files         459      471      +12     
  Lines       29951    30290     +339     
==========================================
+ Hits        25526    26616    +1090     
+ Misses       4425     3674     -751     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9275?dropdown=coverage&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/data/graph_store.py b/torch_geometric/data/graph_store.py
index 4696f3ad6..b5f6e3e23 100644
--- a/torch_geometric/data/graph_store.py
+++ b/torch_geometric/data/graph_store.py
@@ -25,10 +25,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from torch import Tensor
 
+from torch_geometric.index import index2ptr, ptr2index
 from torch_geometric.typing import EdgeTensorType, EdgeType, OptTensor
 from torch_geometric.utils import index_sort
 from torch_geometric.utils.mixin import CastMixin
-from torch_geometric.utils.sparse import index2ptr, ptr2index
 
 # The output of converting between two types in the GraphStore is a Tuple of
 # dictionaries: row, col, and perm. The dictionaries are keyed by the edge
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 8eb174913..dfe72ba7f 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1,5 +1,4 @@
 import functools
-import typing
 from enum import Enum
 from typing import (
     Any,
@@ -11,7 +10,6 @@ from typing import (
     NamedTuple,
     Optional,
     Sequence,
-    Set,
     Tuple,
     Type,
     Union,
@@ -25,22 +23,13 @@ from torch import Tensor
 
 import torch_geometric.typing
 from torch_geometric import is_compiling
-from torch_geometric.typing import SparseTensor
+from torch_geometric.index import index2ptr, ptr2index
+from torch_geometric.typing import INDEX_DTYPES, SparseTensor
 
 aten = torch.ops.aten
 
 HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
 
-if torch_geometric.typing.WITH_PT20:
-    SUPPORTED_DTYPES: Set[torch.dtype] = {
-        torch.int32,
-        torch.int64,
-    }
-elif not typing.TYPE_CHECKING:  # pragma: no cover
-    SUPPORTED_DTYPES: Set[torch.dtype] = {
-        torch.int64,
-    }
-
 ReduceType = Literal['sum', 'mean', 'amin', 'amax', 'add', 'min', 'max']
 PYG_REDUCE: Dict[ReduceType, ReduceType] = {
     'add': 'sum',
@@ -118,16 +107,11 @@ def maybe_sub(
                  for v, o in zip(value, other))
 
 
-def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
-    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
-    return index.repeat_interleave(ptr.diff(), output_size=output_size)
-
-
 def assert_valid_dtype(tensor: Tensor) -> None:
-    if tensor.dtype not in SUPPORTED_DTYPES:
+    if tensor.dtype not in INDEX_DTYPES:
         raise ValueError(f"'EdgeIndex' holds an unsupported data type "
                          f"(got '{tensor.dtype}', but expected one of "
-                         f"{SUPPORTED_DTYPES})")
+                         f"{INDEX_DTYPES})")
 
 
 def assert_two_dimensional(tensor: Tensor) -> None:
@@ -610,11 +594,7 @@ class EdgeIndex(Tensor):
             return self._T_indptr
 
         dim = 0 if self.is_sorted_by_row else 1
-        self._indptr = torch._convert_indices_from_coo_to_csr(
-            self._data[dim],
-            self.get_sparse_size(dim),
-            out_int32=self.dtype != torch.int64,
-        )
+        self._indptr = index2ptr(self._data[dim], self.get_sparse_size(dim))
 
         return self._indptr
 
@@ -655,11 +635,7 @@ class EdgeIndex(Tensor):
         elif self.is_undirected and self._indptr is not None:
             rowptr = self._indptr
         else:
-            rowptr = self._T_indptr = torch._convert_indices_from_coo_to_csr(
-                row,
-                self.get_num_rows(),
-                out_int32=self.dtype != torch.int64,
-            )
+            rowptr = self._T_indptr = index2ptr(row, self.get_num_rows())
 
         return (rowptr, col), perm
 
@@ -679,11 +655,7 @@ class EdgeIndex(Tensor):
         elif self.is_undirected and self._indptr is not None:
             colptr = self._indptr
         else:
-            colptr = self._T_indptr = torch._convert_indices_from_coo_to_csr(
-                col,
-                self.get_num_cols(),
-                out_int32=self.dtype != torch.int64,
-            )
+            colptr = self._T_indptr = index2ptr(col, self.get_num_cols())
 
         return (colptr, row), perm
 
@@ -1282,7 +1254,7 @@ def apply_(
 
     data = fn(tensor._data, *args, **kwargs)
 
-    if data.dtype not in SUPPORTED_DTYPES:
+    if data.dtype not in INDEX_DTYPES:
         return data
 
     if tensor._data.data_ptr() != data.data_ptr():
@@ -1531,7 +1503,7 @@ def _add(
         alpha=alpha,
     )
 
-    if out.dtype not in SUPPORTED_DTYPES:
+    if out.dtype not in INDEX_DTYPES:
         return out
     if out.dim() != 2 or out.size(0) != 2:
         return out
@@ -1628,7 +1600,7 @@ def _sub(
         alpha=alpha,
     )
 
-    if out.dtype not in SUPPORTED_DTYPES:
+    if out.dtype not in INDEX_DTYPES:
         return out
     if out.dim() != 2 or out.size(0) != 2:
         return out
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
new file mode 100644
index 000000000..6c05e330a
--- /dev/null
+++ b/torch_geometric/index.py
@@ -0,0 +1,17 @@
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
+    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
+    return index.repeat_interleave(ptr.diff(), output_size=output_size)
+
+
+def index2ptr(index: Tensor, size: Optional[int] = None) -> Tensor:
+    if size is None:
+        size = int(index.max()) + 1 if index.numel() > 0 else 0
+
+    return torch._convert_indices_from_coo_to_csr(
+        index, size, out_int32=index.dtype != torch.int64)
diff --git a/torch_geometric/loader/cluster.py b/torch_geometric/loader/cluster.py
index a22eb3556..3839b0301 100644
--- a/torch_geometric/loader/cluster.py
+++ b/torch_geometric/loader/cluster.py
@@ -10,10 +10,10 @@ from torch import Tensor
 
 import torch_geometric.typing
 from torch_geometric.data import Data
+from torch_geometric.index import index2ptr, ptr2index
 from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import index_sort, narrow, select, sort_edge_index
 from torch_geometric.utils.map import map_index
-from torch_geometric.utils.sparse import index2ptr, ptr2index
 
 
 @dataclass
diff --git a/torch_geometric/nn/conv/collect.jinja b/torch_geometric/nn/conv/collect.jinja
index f455b577e..b566e7ef6 100644
--- a/torch_geometric/nn/conv/collect.jinja
+++ b/torch_geometric/nn/conv/collect.jinja
@@ -4,8 +4,8 @@ import torch
 from torch import Tensor
 
 from torch_geometric import EdgeIndex
+from torch_geometric.index import ptr2index
 from torch_geometric.utils import is_torch_sparse_tensor
-from torch_geometric.utils.sparse import ptr2index
 from torch_geometric.typing import SparseTensor
 
 
diff --git a/torch_geometric/nn/conv/fused_gat_conv.py b/torch_geometric/nn/conv/fused_gat_conv.py
index 5b8ed378d..a12fc31a0 100644
--- a/torch_geometric/nn/conv/fused_gat_conv.py
+++ b/torch_geometric/nn/conv/fused_gat_conv.py
@@ -3,9 +3,9 @@ from typing import Optional, Tuple
 import torch
 from torch import Tensor
 
+from torch_geometric.index import index2ptr
 from torch_geometric.nn.conv import GATConv
 from torch_geometric.utils import sort_edge_index
-from torch_geometric.utils.sparse import index2ptr
 
 
 class FusedGATConv(GATConv):  # pragma: no cover
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index e282097b0..79f79a1ba 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -20,6 +20,7 @@ from torch import Tensor
 from torch.utils.hooks import RemovableHandle
 
 from torch_geometric import EdgeIndex, is_compiling
+from torch_geometric.index import ptr2index
 from torch_geometric.inspector import Inspector, Signature
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.resolver import aggregation_resolver as aggr_resolver
@@ -30,7 +31,6 @@ from torch_geometric.utils import (
     is_torch_sparse_tensor,
     to_edge_index,
 )
-from torch_geometric.utils.sparse import ptr2index
 
 FUSE_AGGRS = {'add', 'sum', 'mean', 'min', 'max'}
 HookDict = OrderedDict[int, Callable]
diff --git a/torch_geometric/nn/conv/rgcn_conv.py b/torch_geometric/nn/conv/rgcn_conv.py
index 6a042ec0a..f677f4d85 100644
--- a/torch_geometric/nn/conv/rgcn_conv.py
+++ b/torch_geometric/nn/conv/rgcn_conv.py
@@ -7,6 +7,7 @@ from torch.nn import Parameter
 import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric import is_compiling
+from torch_geometric.index import index2ptr
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.inits import glorot, zeros
 from torch_geometric.typing import (
@@ -17,7 +18,6 @@ from torch_geometric.typing import (
     torch_sparse,
 )
 from torch_geometric.utils import index_sort, one_hot, scatter, spmm
-from torch_geometric.utils.sparse import index2ptr
 
 
 def masked_edge_index(edge_index: Adj, edge_mask: Tensor) -> Adj:
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index 138a98eb4..24a173b6a 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -12,10 +12,10 @@ from torch.nn.parameter import Parameter
 import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric import is_compiling
+from torch_geometric.index import index2ptr
 from torch_geometric.nn import inits
 from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import index_sort
-from torch_geometric.utils.sparse import index2ptr
 
 
 def is_uninitialized_parameter(x: Any) -> bool:
diff --git a/torch_geometric/nn/models/metapath2vec.py b/torch_geometric/nn/models/metapath2vec.py
index faac817c8..406d44617 100644
--- a/torch_geometric/nn/models/metapath2vec.py
+++ b/torch_geometric/nn/models/metapath2vec.py
@@ -5,9 +5,9 @@ from torch import Tensor
 from torch.nn import Embedding
 from torch.utils.data import DataLoader
 
+from torch_geometric.index import index2ptr
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 from torch_geometric.utils import sort_edge_index
-from torch_geometric.utils.sparse import index2ptr
 
 EPS = 1e-15
 
diff --git a/torch_geometric/nn/models/node2vec.py b/torch_geometric/nn/models/node2vec.py
index 553b4f637..1b61bc5c7 100644
--- a/torch_geometric/nn/models/node2vec.py
+++ b/torch_geometric/nn/models/node2vec.py
@@ -5,10 +5,10 @@ from torch import Tensor
 from torch.nn import Embedding
 from torch.utils.data import DataLoader
 
+from torch_geometric.index import index2ptr
 from torch_geometric.typing import WITH_PYG_LIB, WITH_TORCH_CLUSTER
 from torch_geometric.utils import sort_edge_index
 from torch_geometric.utils.num_nodes import maybe_num_nodes
-from torch_geometric.utils.sparse import index2ptr
 
 
 class Node2Vec(torch.nn.Module):
diff --git a/torch_geometric/sampler/utils.py b/torch_geometric/sampler/utils.py
index 584eec187..c0eda2a3b 100644
--- a/torch_geometric/sampler/utils.py
+++ b/torch_geometric/sampler/utils.py
@@ -5,9 +5,9 @@ from torch import Tensor
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.storage import EdgeStorage
+from torch_geometric.index import index2ptr
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 from torch_geometric.utils import coalesce, index_sort, lexsort
-from torch_geometric.utils.sparse import index2ptr
 
 # Edge Layout Conversion ######################################################
 
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 049c86bb8..3cf619447 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -1,8 +1,9 @@
 import inspect
 import os
 import sys
+import typing
 import warnings
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -21,6 +22,16 @@ NO_MKL = 'USE_MKL=OFF' in torch.__config__.show() or WITH_WINDOWS
 
 MAX_INT64 = torch.iinfo(torch.int64).max
 
+if WITH_PT20:
+    INDEX_DTYPES: Set[torch.dtype] = {
+        torch.int32,
+        torch.int64,
+    }
+elif not typing.TYPE_CHECKING:  # pragma: no cover
+    INDEX_DTYPES: Set[torch.dtype] = {
+        torch.int64,
+    }
+
 if not hasattr(torch, 'sparse_csc'):
     torch.sparse_csc = torch.sparse_coo
 
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index 6800d7cae..d9ca98e07 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric.index import index2ptr, ptr2index
 from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import coalesce, cumsum
 
@@ -483,19 +484,6 @@ def set_sparse_value(adj: Tensor, value: Tensor) -> Tensor:
     raise ValueError(f"Unexpected sparse tensor layout (got '{adj.layout}')")
 
 
-def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
-    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
-    return index.repeat_interleave(ptr.diff(), output_size=output_size)
-
-
-def index2ptr(index: Tensor, size: Optional[int] = None) -> Tensor:
-    if size is None:
-        size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-    return torch._convert_indices_from_coo_to_csr(
-        index, size, out_int32=index.dtype == torch.int32)
-
-
 def cat_coo(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
     assert dim in {0, 1, (0, 1)}
     assert tensors[0].layout == torch.sparse_coo


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 1d4abe629..37fe4609c 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -9,7 +9,6 @@ from torch import Tensor, tensor
 import torch_geometric
 from torch_geometric import EdgeIndex
 from torch_geometric.edge_index import (
-    SUPPORTED_DTYPES,
     ReduceType,
     SortReturnType,
     _scatter_spmm,
@@ -25,10 +24,10 @@ from torch_geometric.testing import (
     withoutExtensions,
     withPackage,
 )
-from torch_geometric.typing import SparseTensor
+from torch_geometric.typing import INDEX_DTYPES, SparseTensor
 from torch_geometric.utils import scatter
 
-DTYPES = [pytest.param(dtype, id=str(dtype)[6:]) for dtype in SUPPORTED_DTYPES]
+DTYPES = [pytest.param(dtype, id=str(dtype)[6:]) for dtype in INDEX_DTYPES]
 IS_UNDIRECTED = [
     pytest.param(False, id='directed'),
     pytest.param(True, id='undirected'),
diff --git a/test/utils/test_segment.py b/test/utils/test_segment.py
index 541b2c742..6eaa2ecc3 100644
--- a/test/utils/test_segment.py
+++ b/test/utils/test_segment.py
@@ -4,10 +4,10 @@ import pytest
 import torch
 
 import torch_geometric.typing
+from torch_geometric.index import index2ptr
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withoutExtensions
 from torch_geometric.utils import scatter, segment
-from torch_geometric.utils.sparse import index2ptr
 
 
 @withCUDA