Repository: pyg-team/pytorch_geometric
Pull request #8681 (2023-12-29 09:07:36)

Drop `torch-sparse` requirement in `AddMetaPaths`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`e3cea7c`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/e3cea7ca1c6a6f3bd0179f9816f6330d2fc1ea8b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 90.13% compared to head [(`5e67e03`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.47%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?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/transforms/add\_metapaths.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvYWRkX21ldGFwYXRocy5weQ==) | 97.36% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?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    #8681      +/-   ##
==========================================
- Coverage   90.13%   89.47%   -0.66%     
==========================================
  Files         481      481              
  Lines       30988    30989       +1     
==========================================
- Hits        27930    27728     -202     
- Misses       3058     3261     +203     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8681?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 67653db20..0fab30d29 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -7,6 +7,7 @@ from typing import (
     Dict,
     List,
     Literal,
+    NamedTuple,
     Optional,
     Sequence,
     Set,
@@ -294,6 +295,15 @@ class EdgeIndex(Tensor):
             out._T_indptr = data._T_indptr
             out._value = out._value
 
+            # Reset metadata if cache is invalidated:
+            num_rows = sparse_size[0]
+            if num_rows is not None and num_rows != data.sparse_size(0):
+                out._indptr = None
+
+            num_cols = sparse_size[1]
+            if num_cols is not None and num_cols != data.sparse_size(1):
+                out._T_indptr = None
+
         return out
 
     # Validation ##############################################################
@@ -595,7 +605,7 @@ class EdgeIndex(Tensor):
         self,
         sort_order: Union[str, SortOrder],
         stable: bool = False,
-    ) -> torch.return_types.sort:
+    ) -> 'SortReturnType':
         r"""Sorts the elements by row or column indices.
 
         Args:
@@ -610,10 +620,7 @@ class EdgeIndex(Tensor):
         sort_order = SortOrder(sort_order)
 
         if self._sort_order == sort_order:  # Nothing to do.
-            return torch.return_types.sort([  # type: ignore
-                self,
-                slice(None, None, None),
-            ])
+            return SortReturnType(self, slice(None, None, None))
 
         if self.is_sorted:
             (row, col), perm = self._sort_by_transpose()
@@ -630,20 +637,20 @@ class EdgeIndex(Tensor):
 
         out = self.__class__(edge_index)
 
-        # We can mostly inherit metadata and cache:
+        # We can inherit metadata and (mostly) cache:
         out._sparse_size = self.sparse_size()
         out._sort_order = sort_order
         out._is_undirected = self.is_undirected
 
-        out._indptr = self._T_indptr
-        out._T_indptr = self._indptr
+        out._indptr = self._indptr
+        out._T_indptr = self._T_indptr
 
         # NOTE We cannot copy CSR<>CSC permutations since we don't require that
         # local neighborhoods are sorted, and thus they may run out of sync.
 
         out._value = self._value
 
-        return torch.return_types.sort([out, perm])  # type: ignore
+        return SortReturnType(out, perm)
 
     def to_dense(  # type: ignore
         self,
@@ -653,6 +660,13 @@ class EdgeIndex(Tensor):
     ) -> Tensor:
         r"""Converts :class:`EdgeIndex` into a dense :class:`torch.Tensor`.
 
+        .. warning::
+
+            In case of duplicated edges, the behavior is non-deterministic (one
+            of the values from :obj:`value` will be picked arbitrarily). For
+            deterministic behavior, consider calling
+            :meth:`~torch_geometric.utils.coalesce` beforehand.
+
         Args:
             value (torch.Tensor, optional): The values for non-zero elements.
                 If not specified, non-zero elements will be assigned a value of
@@ -662,8 +676,6 @@ class EdgeIndex(Tensor):
             dtype (torch.dtype, optional): The data type of the returned
                 tensor. (default: :obj:`None`)
         """
-        # TODO Respect duplicated edges.
-
         dtype = value.dtype if value is not None else dtype
 
         size = self.get_sparse_size()
@@ -906,6 +918,11 @@ class EdgeIndex(Tensor):
         return Tensor.__torch_function__(func, _types, args, kwargs)
 
 
+class SortReturnType(NamedTuple):
+    values: EdgeIndex
+    indices: Union[Tensor, slice]
+
+
 def apply_(
     tensor: EdgeIndex,
     fn: Callable,
diff --git a/torch_geometric/transforms/add_metapaths.py b/torch_geometric/transforms/add_metapaths.py
index ade571eca..711f931c8 100644
--- a/torch_geometric/transforms/add_metapaths.py
+++ b/torch_geometric/transforms/add_metapaths.py
@@ -1,13 +1,14 @@
 import warnings
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, cast
 
 import torch
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data import HeteroData
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
-from torch_geometric.typing import EdgeType, SparseTensor
+from torch_geometric.typing import EdgeType
 from torch_geometric.utils import coalesce, degree
 
 
@@ -90,10 +91,10 @@ class AddMetaPaths(BaseTransform):
         max_sample (int, optional): If set, will sample at maximum
             :obj:`max_sample` neighbors within metapaths. Useful in order to
             tackle very dense metapath edges. (default: :obj:`None`)
-        weighted (bool, optional): If set to :obj:`True` compute weights for
-            each metapath edge and store them in :obj:`edge_weight`. The weight
-            of each metapath edge is computed as the number of metapaths from
-            the start to the end of the metapath edge.
+        weighted (bool, optional): If set to :obj:`True`, computes weights for
+            each metapath edge and stores them in :obj:`edge_weight`. The
+            weight of each metapath edge is computed as the number of metapaths
+            from the start to the end of the metapath edge.
             (default :obj:`False`)
     """
     def __init__(
@@ -130,38 +131,34 @@ class AddMetaPaths(BaseTransform):
         self.weighted = weighted
 
     def forward(self, data: HeteroData) -> HeteroData:
-        edge_types = data.edge_types  # save original edge types
+        edge_types = data.edge_types  # Save original edge types.
         data.metapath_dict = {}
 
         for j, metapath in enumerate(self.metapaths):
             for edge_type in metapath:
-                assert data._to_canonical(
-                    edge_type) in edge_types, f"'{edge_type}' not present"
+                assert data._to_canonical(edge_type) in edge_types
 
             edge_type = metapath[0]
-            edge_weight = self._get_edge_weight(data, edge_type)
-            # TODO Drop `torch-sparse` requirement.
-            adj1 = SparseTensor.from_edge_index(
-                edge_index=data[edge_type].edge_index,
-                sparse_sizes=data[edge_type].size(), edge_attr=edge_weight)
+            edge_index, edge_weight = self._edge_index(data, edge_type)
 
             if self.max_sample is not None:
-                adj1 = self.sample_adj(adj1)
+                edge_index, edge_weight = self._sample(edge_index, edge_weight)
 
             for i, edge_type in enumerate(metapath[1:]):
-                edge_weight = self._get_edge_weight(data, edge_type)
-                adj2 = SparseTensor.from_edge_index(
-                    edge_index=data[edge_type].edge_index,
-                    sparse_sizes=data[edge_type].size(), edge_attr=edge_weight)
+                edge_index2, edge_weight2 = self._edge_index(data, edge_type)
+
+                edge_index, edge_weight = edge_index.matmul(
+                    edge_index2, edge_weight, edge_weight2)
 
-                adj1 = adj1 @ adj2
+                if not self.weighted:
+                    edge_weight = None
 
                 if self.max_sample is not None:
-                    adj1 = self.sample_adj(adj1)
+                    edge_index, edge_weight = self._sample(
+                        edge_index, edge_weight)
 
-            row, col, edge_weight = adj1.coo()
             new_edge_type = (metapath[0][0], f'metapath_{j}', metapath[-1][-1])
-            data[new_edge_type].edge_index = torch.vstack([row, col])
+            data[new_edge_type].edge_index = edge_index
             if self.weighted:
                 data[new_edge_type].edge_weight = edge_weight
             data.metapath_dict[new_edge_type] = metapath
@@ -171,25 +168,44 @@ class AddMetaPaths(BaseTransform):
 
         return data
 
-    def sample_adj(self, adj: SparseTensor) -> SparseTensor:
-        row, col, _ = adj.coo()
-        deg = degree(row, num_nodes=adj.size(0))
-        prob = (self.max_sample * (1. / deg))[row]
+    def _edge_index(
+        self,
+        data: HeteroData,
+        edge_type: EdgeType,
+    ) -> Tuple[EdgeIndex, Optional[Tensor]]:
+
+        edge_index = EdgeIndex(
+            data[edge_type].edge_index,
+            sparse_size=data[edge_type].size(),
+        )
+        edge_index, perm = edge_index.sort_by('row')
+
+        if not self.weighted:
+            return edge_index, None
+
+        edge_weight = data[edge_type].get('edge_weight')
+        if edge_weight is not None:
+            assert edge_weight.dim() == 1
+            edge_weight = edge_weight[perm]
+
+        return edge_index, edge_weight
+
+    def _sample(
+        self,
+        edge_index: EdgeIndex,
+        edge_weight: Optional[Tensor],
+    ) -> Tuple[EdgeIndex, Optional[Tensor]]:
+
+        deg = degree(edge_index[0], num_nodes=edge_index.get_sparse_size(0))
+        prob = (self.max_sample * (1. / deg))[edge_index[0]]
         mask = torch.rand_like(prob) < prob
-        return adj.masked_select_nnz(mask, layout='coo')
-
-    def _get_edge_weight(self, data: HeteroData,
-                         edge_type: EdgeType) -> torch.Tensor:
-        if self.weighted:
-            edge_weight = data[edge_type].get('edge_weight', None)
-            if edge_weight is None:
-                edge_weight = torch.ones(
-                    data[edge_type].num_edges,
-                    device=data[edge_type].edge_index.device)
-            assert edge_weight.ndim == 1
-        else:
-            edge_weight = None
-        return edge_weight
+
+        edge_index = cast(EdgeIndex, edge_index[:, mask])
+        assert isinstance(edge_index, EdgeIndex)
+        if edge_weight is not None:
+            edge_weight = edge_weight[mask]
+
+        return edge_index, edge_weight
 
 
 @functional_transform('add_random_metapaths')
@@ -261,10 +277,11 @@ class AddRandomMetaPaths(BaseTransform):
                 self.walks_per_node[j])
 
             for i, edge_type in enumerate(metapath):
-                edge_index = data[edge_type].edge_index
-                adj = SparseTensor(row=edge_index[0], col=edge_index[1],
-                                   sparse_sizes=data[edge_type].size())
-                col, mask = self.sample(adj, start)
+                edge_index = EdgeIndex(
+                    data[edge_type].edge_index,
+                    sparse_size=data[edge_type].size(),
+                )
+                col, mask = self.sample(edge_index, start)
                 row, col = row[mask], col[mask]
                 start = col
 
@@ -278,18 +295,23 @@ class AddRandomMetaPaths(BaseTransform):
         return data
 
     @staticmethod
-    def sample(adj: SparseTensor, subset: Tensor) -> Tuple[Tensor, Tensor]:
-        """Sample a neighbor form :obj:`adj` for each node in :obj:`subset`."""
-        rowcount = adj.storage.rowcount()[subset]
+    def sample(edge_index: EdgeIndex, subset: Tensor) -> Tuple[Tensor, Tensor]:
+        """Sample neighbors from :obj:`edge_index` for each node in
+        :obj:`subset`.
+        """
+        edge_index, _ = edge_index.sort_by('row')
+        rowptr = edge_index.get_indptr()
+        rowcount = rowptr.diff()[subset]
+
         mask = rowcount > 0
         offset = torch.zeros_like(subset)
-        offset[mask] = adj.storage.rowptr()[subset[mask]]
+        offset[mask] = rowptr[subset[mask]]
 
         rand = torch.rand((rowcount.size(0), 1), device=subset.device)
         rand.mul_(rowcount.to(rand.dtype).view(-1, 1))
         rand = rand.to(torch.long)
         rand.add_(offset.view(-1, 1))
-        col = adj.storage.col()[rand].squeeze()
+        col = edge_index[1][rand].squeeze()
         return col, mask
 
     def __repr__(self) -> str:
@@ -313,7 +335,7 @@ def postprocess(
             else:
                 del data[i]
 
-    # remove nodes not connected by any edge type.
+    # Remove nodes not connected by any edge type:
     if drop_unconnected_node_types:
         new_edge_types = data.edge_types
         node_types = data.node_types


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 68ee5a1d4..213f910c7 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -11,6 +11,7 @@ from torch_geometric import EdgeIndex
 from torch_geometric.edge_index import (
     SUPPORTED_DTYPES,
     ReduceType,
+    SortReturnType,
     _scatter_spmm,
     _torch_sparse_spmm,
     _TorchSPMM,
@@ -311,7 +312,7 @@ def test_sort_by(dtype, device, is_undirected):
     kwargs = dict(dtype=dtype, device=device, is_undirected=is_undirected)
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row', **kwargs)
     out = adj.sort_by('row')
-    assert isinstance(out, torch.return_types.sort)
+    assert isinstance(out, SortReturnType)
     assert isinstance(out.values, EdgeIndex)
     assert not isinstance(out.indices, EdgeIndex)
     assert out.values.equal(adj)
@@ -319,7 +320,7 @@ def test_sort_by(dtype, device, is_undirected):
 
     adj = EdgeIndex([[0, 1, 2, 1], [1, 0, 1, 2]], **kwargs)
     out = adj.sort_by('row')
-    assert isinstance(out, torch.return_types.sort)
+    assert isinstance(out, SortReturnType)
     assert isinstance(out.values, EdgeIndex)
     assert not isinstance(out.indices, EdgeIndex)
     assert out.values[0].equal(tensor([0, 1, 1, 2], device=device))
diff --git a/test/transforms/test_add_metapaths.py b/test/transforms/test_add_metapaths.py
index 994ad3527..497e9e4e9 100644
--- a/test/transforms/test_add_metapaths.py
+++ b/test/transforms/test_add_metapaths.py
@@ -2,7 +2,6 @@ import torch
 from torch import tensor
 
 from torch_geometric.data import HeteroData
-from torch_geometric.testing import withPackage
 from torch_geometric.transforms import AddMetaPaths, AddRandomMetaPaths
 from torch_geometric.utils import coalesce
 
@@ -20,8 +19,7 @@ def generate_data() -> HeteroData:
     return data
 
 
-@withPackage('torch_sparse')
-def test_add_metapaths():
+def test_add_metapaths() -> None:
     data = generate_data()
     # Test transform options:
     metapaths = [[('p', 'c'), ('c', 'p')]]
@@ -74,8 +72,7 @@ def test_add_metapaths():
     assert list(meta.metapath_dict.keys()) == new_edge_types
 
 
-@withPackage('torch_sparse')
-def test_add_metapaths_max_sample():
+def test_add_metapaths_max_sample() -> None:
     torch.manual_seed(12345)
 
     data = generate_data()
@@ -87,8 +84,7 @@ def test_add_metapaths_max_sample():
     assert meta['metapath_0'].edge_index.size(1) < 9
 
 
-@withPackage('torch_sparse')
-def test_add_weighted_metapaths():
+def test_add_weighted_metapaths() -> None:
     torch.manual_seed(12345)
 
     data = HeteroData()
@@ -146,8 +142,7 @@ def test_add_weighted_metapaths():
     assert edge_weight.tolist() == [1, 2, 2, 4]
 
 
-@withPackage('torch_sparse')
-def test_add_random_metapaths():
+def test_add_random_metapaths() -> None:
     data = generate_data()
 
     # Test transform options: