Repository: pyg-team/pytorch_geometric
Pull request #7051 (2023-03-27 15:47:08)

Drop `torch_sparse` dependency in tests (5/n)


Patch:
diff --git a/torch_geometric/nn/aggr/multi.py b/torch_geometric/nn/aggr/multi.py
index b4855b1e1..6883f935d 100644
--- a/torch_geometric/nn/aggr/multi.py
+++ b/torch_geometric/nn/aggr/multi.py
@@ -182,7 +182,7 @@ class MultiAggregation(Aggregation):
 
         if hasattr(self, 'multihead_attn'):
             x = torch.stack(
-                [head(x) for x, head in zip(inputs, self.lin_heads)],
+                [head(inputs[i]) for i, head in enumerate(self.lin_heads)],
                 dim=0,
             )
             attn_out, _ = self.multihead_attn(x, x, x)


Test patch:
diff --git a/.github/workflows/latest_testing.yml b/.github/workflows/latest_testing.yml
index f39918ce1..4f1cf1f9a 100644
--- a/.github/workflows/latest_testing.yml
+++ b/.github/workflows/latest_testing.yml
@@ -60,7 +60,7 @@ jobs:
           pytest test/transforms/
           pytest test/utils/
           pytest test/visualization/
-          # pytest test/nn/aggr
+          pytest test/nn/aggr
           # pytest test/nn/conv
           pytest test/nn/dense
           pytest test/nn/functional
diff --git a/test/nn/aggr/test_attention.py b/test/nn/aggr/test_attention.py
index 2bb620d78..012e1fdb1 100644
--- a/test/nn/aggr/test_attention.py
+++ b/test/nn/aggr/test_attention.py
@@ -1,5 +1,7 @@
+import pytest
 import torch
 
+import torch_geometric.typing
 from torch_geometric.nn import MLP
 from torch_geometric.nn.aggr import AttentionalAggregation
 
@@ -18,4 +20,9 @@ def test_attentional_aggregation():
 
     out = aggr(x, index)
     assert out.size() == (3, channels)
-    torch.allclose(aggr(x, ptr=ptr, dim_size=3), out)
+
+    if not torch_geometric.typing.WITH_TORCH_SCATTER:
+        with pytest.raises(ImportError, match="'segment' requires"):
+            aggr(x, ptr=ptr)
+    else:
+        assert torch.allclose(out, aggr(x, ptr=ptr))
diff --git a/test/nn/aggr/test_basic.py b/test/nn/aggr/test_basic.py
index 5370ad82d..45d04dc3f 100644
--- a/test/nn/aggr/test_basic.py
+++ b/test/nn/aggr/test_basic.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+import torch_geometric.typing
 from torch_geometric.nn import (
     MaxAggregation,
     MeanAggregation,
@@ -54,6 +55,9 @@ def test_basic_aggregation(Aggregation):
     if isinstance(aggr, MulAggregation):
         with pytest.raises(NotImplementedError, match="requires 'index'"):
             aggr(x, ptr=ptr)
+    elif not torch_geometric.typing.WITH_TORCH_SCATTER:
+        with pytest.raises(ImportError, match="'segment' requires"):
+            aggr(x, ptr=ptr)
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
 
@@ -75,7 +79,7 @@ def test_var_aggregation():
     PowerMeanAggregation,
 ])
 @pytest.mark.parametrize('learn', [True, False])
-def test_gen_aggregation(Aggregation, learn):
+def test_learnable_aggregation(Aggregation, learn):
     x = torch.randn(6, 16)
     index = torch.tensor([0, 0, 1, 1, 1, 2])
     ptr = torch.tensor([0, 2, 5, 6])
@@ -85,7 +89,12 @@ def test_gen_aggregation(Aggregation, learn):
 
     out = aggr(x, index)
     assert out.size() == (3, x.size(1))
-    assert torch.allclose(out, aggr(x, ptr=ptr))
+
+    if not torch_geometric.typing.WITH_TORCH_SCATTER:
+        with pytest.raises(ImportError, match="'segment' requires"):
+            aggr(x, ptr=ptr)
+    else:
+        assert torch.allclose(out, aggr(x, ptr=ptr))
 
     if learn:
         out.mean().backward()
@@ -107,7 +116,12 @@ def test_learnable_channels_aggregation(Aggregation):
 
     out = aggr(x, index)
     assert out.size() == (3, x.size(1))
-    assert torch.allclose(out, aggr(x, ptr=ptr))
+
+    if not torch_geometric.typing.WITH_TORCH_SCATTER:
+        with pytest.raises(ImportError, match="'segment' requires"):
+            aggr(x, ptr=ptr)
+    else:
+        assert torch.allclose(out, aggr(x, ptr=ptr))
 
     out.mean().backward()
     for param in aggr.parameters():
diff --git a/test/nn/aggr/test_multi.py b/test/nn/aggr/test_multi.py
index ccfa70dbf..222696b62 100644
--- a/test/nn/aggr/test_multi.py
+++ b/test/nn/aggr/test_multi.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+import torch_geometric.typing
 from torch_geometric.nn import MultiAggregation
 
 
@@ -36,7 +37,13 @@ def test_multi_aggr(multi_aggr_tuple):
                          f"], mode={aggr_kwargs['mode']})")
 
     out = aggr(x, index)
-    assert torch.allclose(out, aggr(x, ptr=ptr))
     assert out.size() == (4, expand * x.size(1))
 
-    # TODO test JIT support
+    if not torch_geometric.typing.WITH_TORCH_SCATTER:
+        with pytest.raises(ImportError, match="'segment' requires"):
+            aggr(x, ptr=ptr)
+    else:
+        assert torch.allclose(out, aggr(x, ptr=ptr))
+
+    jit = torch.jit.script(aggr)
+    assert torch.allclose(out, jit(x, index))