Repository: pyg-team/pytorch_geometric
Pull request #8481 (2023-11-29 16:25:43)

Test `EdgeIndex.matmul` backward

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?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 [(`1ba2743`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/1ba27434f86429b178adbc1724ff03bc42f289c4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.69% compared to head [(`5476280`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.37%.
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?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/data/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 85.71% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [...rch\_geometric/distributed/dist\_neighbor\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2Rpc3RyaWJ1dGVkL2Rpc3RfbmVpZ2hib3Jfc2FtcGxlci5weQ==) | 93.75% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?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    #8481      +/-   ##
==========================================
+ Coverage   87.69%   88.37%   +0.68%     
==========================================
  Files         478      478              
  Lines       29400    29405       +5     
==========================================
+ Hits        25781    25986     +205     
+ Misses       3619     3419     -200     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8481?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/edge_index.py b/torch_geometric/data/edge_index.py
index 4e6e4fdb0..ca9cbb4c0 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -629,7 +629,7 @@ def to_sparse_csr(tensor: EdgeIndex, value: Optional[Tensor] = None) -> Tensor:
     return torch.sparse_csr_tensor(
         crow_indices=tensor.get_rowptr(),
         col_indices=tensor[1],
-        values=_get_value(tensor.size(1), device=tensor.device),
+        values=value,
         size=tensor.get_sparse_size(),
         device=tensor.device,
     )
@@ -695,8 +695,8 @@ else:
 def matmul(
     input: EdgeIndex,
     other: Union[Tensor, EdgeIndex],
-    input_weight: Optional[Tensor] = None,
-    other_weight: Optional[Tensor] = None,
+    input_value: Optional[Tensor] = None,
+    other_value: Optional[Tensor] = None,
     reduce: Literal['sum'] = 'sum',
 ) -> Union[Tensor, Tuple[EdgeIndex, Tensor]]:
 
@@ -705,15 +705,18 @@ def matmul(
     # TODO Utilize available `CSC` representation for faster backward passes.
 
     if input._sort_order == SortOrder.COL:
-        input = to_sparse_csc(input, other_weight)
+        input = to_sparse_csc(input, input_value)
     else:
-        input = to_sparse_csr(input, other_weight)
+        input = to_sparse_csr(input, input_value)
 
     if isinstance(other, EdgeIndex):
         if other._sort_order == SortOrder.COL:
-            other = to_sparse_csc(other, other_weight)
+            other = to_sparse_csc(other, other_value)
         else:
-            other = to_sparse_csr(other, other_weight)
+            other = to_sparse_csr(other, other_value)
+
+    elif other_value is not None:
+        raise ValueError("'other_value' not supported for sparse-dense "
+                         "matrix multiplication")
 
-    return Tensor.__torch_function__(  #
-        Tensor.matmul, (Tensor, ), (input, other))
+    return torch.matmul(input, other)


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 3ce830b6e..141b6696b 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -6,7 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric
-from torch_geometric.data.edge_index import EdgeIndex, to_dense
+from torch_geometric.data.edge_index import EdgeIndex, matmul, to_dense
 from torch_geometric.testing import (
     disableExtensions,
     onlyCUDA,
@@ -323,6 +323,36 @@ def test_matmul():
     assert torch.allclose(out.to_dense(), adj2.to_dense() @ adj2.to_dense())
 
 
+def test_matmul_input_value():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    x = torch.randn(3, 1)
+    value = torch.randn(4)
+
+    out = matmul(adj, x, input_value=value)
+    assert torch.allclose(out, to_dense(adj, value=value) @ x)
+
+
+def test_matmul_grad():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    x1 = torch.randn(3, 1, requires_grad=True)
+    value = torch.randn(4, requires_grad=True)
+
+    out = matmul(adj, x1, input_value=value)
+    grad_out = torch.randn_like(out)
+    out.backward(grad_out)
+
+    x2 = x1.detach().requires_grad_()
+    dense_adj = to_dense(adj, value=value).detach().requires_grad_()
+    out = dense_adj @ x2
+    out.backward(grad_out)
+
+    assert torch.allclose(x1.grad, x2.grad)
+    if torch_geometric.typing.WITH_PT21:  # TODO Investigate.
+        assert torch.allclose(value.grad, dense_adj.grad[adj[0], adj[1]])
+
+
 def test_save_and_load(tmp_path):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.fill_cache()