Repository: pyg-team/pytorch_geometric
Pull request #8535 (2023-12-05 09:54:34)

Add `EdgeIndex.spspmm` tests


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 64ad3ec46..d05f96200 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -1247,9 +1247,11 @@ def matmul(
                              "matrix multiplication")
         return _spmm(input, other, input_value, reduce, transpose)
 
-    if transpose:
-        raise NotImplementedError("`transpose=True` not yet supported for "
-                                  "sparse-sparse matrix multiplication")
+    if reduce not in ['sum', 'add']:
+        raise NotImplementedError(f"`reduce='{reduce}'` not yet supported for "
+                                  f"sparse-sparse matrix multiplication")
+
+    transpose &= not input.is_undirected or input_value is not None
 
     if torch_geometric.typing.WITH_WINDOWS:
         input = input.to_sparse_coo(input_value)
@@ -1258,6 +1260,9 @@ def matmul(
     else:
         input = input.to_sparse_csr(input_value)
 
+    if transpose:
+        input = input.t()
+
     if torch_geometric.typing.WITH_WINDOWS:
         other = other.to_sparse_coo(input_value)
     elif other.is_sorted_by_col:


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index d72c9997b..6de23c8ef 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -763,31 +763,34 @@ def test_spmm(without_extensions, device, reduce, transpose, is_undirected):
     assert value1.grad.allclose(value2.grad)
 
 
-def test_spspmm():
-    adj1 = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
-    adj1_dense = adj1.to_dense()
-    adj2 = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col')
-    adj2_dense = adj2.to_dense()
-
-    out, value = adj1 @ adj1
-    assert isinstance(out, EdgeIndex)
-    assert out.is_sorted_by_row
-    assert out._sparse_size == (3, 3)
-    if not torch_geometric.typing.WITH_WINDOWS:
-        assert out._indptr is not None
-    assert torch.allclose(out.to_dense(value), adj1_dense @ adj1_dense)
-
-    out, value = adj1 @ adj2
-    assert isinstance(out, EdgeIndex)
-    assert torch.allclose(out.to_dense(value), adj1_dense @ adj2_dense)
+@withCUDA
+@pytest.mark.parametrize('reduce', ReduceType.__args__)
+@pytest.mark.parametrize('transpose', TRANSPOSE)
+@pytest.mark.parametrize('is_undirected', IS_UNDIRECTED)
+def test_spspmm(device, reduce, transpose, is_undirected):
+    if is_undirected:
+        kwargs = dict(device=device, sort_order='row', is_undirected=True)
+        adj1 = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], **kwargs)
+    else:
+        kwargs = dict(device=device, sort_order='row')
+        adj1 = EdgeIndex([[0, 1, 1, 2], [2, 0, 1, 2]], **kwargs)
 
-    out, value = adj2 @ adj1
-    assert isinstance(out, EdgeIndex)
-    assert torch.allclose(out.to_dense(value), adj2_dense @ adj1_dense)
+    adj1_dense = adj1.to_dense().t() if transpose else adj1.to_dense()
+    adj2 = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col',
+                     device=device)
+    adj2_dense = adj2.to_dense()
 
-    out, value = adj2 @ adj2
-    assert isinstance(out, EdgeIndex)
-    assert torch.allclose(out.to_dense(value), adj2_dense @ adj2_dense)
+    if reduce in ['sum', 'add']:
+        out, value = adj1.matmul(adj2, reduce=reduce, transpose=transpose)
+        assert isinstance(out, EdgeIndex)
+        assert out.is_sorted_by_row
+        assert out._sparse_size == (3, 3)
+        if not torch_geometric.typing.WITH_WINDOWS:
+            assert out._indptr is not None
+        assert torch.allclose(out.to_dense(value), adj1_dense @ adj2_dense)
+    else:
+        with pytest.raises(NotImplementedError, match="not yet supported"):
+            adj1.matmul(adj2, reduce=reduce, transpose=transpose)
 
 
 @withCUDA