Repository: pyg-team/pytorch_geometric
Pull request #9026 (2024-03-07 06:27:17)

Support `EdgeIndex` in `spmm`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9026?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:
> Project coverage is 89.24%. Comparing base [(`af0f5f4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/af0f5f44d4dfa8accd898df3e8523d06b67940b0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`89dbf18`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9026?src=pr&el=desc&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    #9026      +/-   ##
==========================================
- Coverage   89.86%   89.24%   -0.63%     
==========================================
  Files         468      468              
  Lines       30039    30038       -1     
==========================================
- Hits        26995    26807     -188     
- Misses       3044     3231     +187     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9026?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/CHANGELOG.md b/CHANGELOG.md
index 56d72ad53..e74fd5bb6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `EdgeIndex` in `spmm` ([#9026](https://github.com/pyg-team/pytorch_geometric/pull/9026))
 - Added option to pre-allocate memory in GPU-based `ApproxKNN` ([#9046](https://github.com/pyg-team/pytorch_geometric/pull/9046))
 - Added support for `EdgeIndex` in `MessagePassing` ([#9007](https://github.com/pyg-team/pytorch_geometric/pull/9007))
 - Added support for `torch.compile` in combination with `EdgeIndex` ([#9007](https://github.com/pyg-team/pytorch_geometric/pull/9007))
diff --git a/torch_geometric/utils/_spmm.py b/torch_geometric/utils/_spmm.py
index ec89adda8..a0bc34cd7 100644
--- a/torch_geometric/utils/_spmm.py
+++ b/torch_geometric/utils/_spmm.py
@@ -4,6 +4,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.typing import Adj, SparseTensor, torch_sparse
 from torch_geometric.utils import is_torch_sparse_tensor, scatter
 
@@ -16,9 +17,11 @@ def spmm(
     r"""Matrix product of sparse matrix with dense matrix.
 
     Args:
-        src (torch.Tensor or torch_sparse.SparseTensor): The input sparse
-            matrix, either a :pyg:`PyG` :class:`torch_sparse.SparseTensor` or a
-            :pytorch:`PyTorch` :class:`torch.sparse.Tensor`.
+        src (torch.Tensor or torch_sparse.SparseTensor or EdgeIndex):
+            The input sparse matrix which can be a
+            :pyg:`PyG` :class:`torch_sparse.SparseTensor`,
+            a :pytorch:`PyTorch` :class:`torch.sparse.Tensor` or
+            a :pyg:`PyG` :class:`EdgeIndex`.
         other (torch.Tensor): The input dense matrix.
         reduce (str, optional): The reduce operation to use
             (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`).
@@ -31,6 +34,9 @@ def spmm(
     if reduce not in ['sum', 'mean', 'min', 'max']:
         raise ValueError(f"`reduce` argument '{reduce}' not supported")
 
+    if not torch.jit.is_scripting() and isinstance(src, EdgeIndex):
+        return src.matmul(other=other, reduce=reduce)  # type: ignore
+
     if isinstance(src, SparseTensor):
         if src.nnz() == 0:
             return other.new_zeros(src.size(0), other.size(1))


Test patch:
diff --git a/test/utils/test_spmm.py b/test/utils/test_spmm.py
index a856a3cb3..76cc8cd94 100644
--- a/test/utils/test_spmm.py
+++ b/test/utils/test_spmm.py
@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric import EdgeIndex
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
 from torch_geometric.typing import SparseTensor
@@ -106,6 +107,24 @@ def test_spmm_jit(reduce):
         assert torch.allclose(out2, out3, atol=1e-6)
 
 
+@withCUDA
+@withPackage('torch>=2.0.0')
+@pytest.mark.parametrize('reduce', ['sum', 'mean', 'min', 'max'])
+def test_spmm_edge_index(device, reduce):
+    src = EdgeIndex(
+        [[0, 1, 1, 2], [1, 0, 2, 1]],
+        sparse_size=(4, 3),
+        sort_order='row',
+        device=device,
+    )
+    other = torch.rand(3, 4, device=device)
+    out = spmm(src, other, reduce=reduce)
+    assert out.size() == (4, 4)
+
+    out2 = spmm(src.to_sparse_coo(), other, reduce=reduce)
+    assert torch.allclose(out, out2)
+
+
 if __name__ == '__main__':
     import argparse