Repository: pyg-team/pytorch_geometric
Pull request #7324 (2023-05-08 13:48:52)

Support `MultiAggregation` in `aggregation_resolver`


Patch:
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index 95339bf18..747f041c5 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -21,7 +21,7 @@ import torch
 from torch import Tensor
 from torch.utils.hooks import RemovableHandle
 
-from torch_geometric.nn.aggr import Aggregation, MultiAggregation
+from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.conv.utils.inspector import (
     Inspector,
     func_body_repr,
@@ -130,18 +130,12 @@ class MessagePassing(torch.nn.Module):
 
         if aggr is None:
             self.aggr = None
-            self.aggr_module = None
         elif isinstance(aggr, (str, Aggregation)):
             self.aggr = str(aggr)
-            self.aggr_module = aggr_resolver(aggr, **(aggr_kwargs or {}))
         elif isinstance(aggr, (tuple, list)):
             self.aggr = [str(x) for x in aggr]
-            self.aggr_module = MultiAggregation(aggr, **(aggr_kwargs or {}))
-        else:
-            raise ValueError(
-                f"Only strings, list, tuples and instances of"
-                f"`torch_geometric.nn.aggr.Aggregation` are "
-                f"valid aggregation schemes (got '{type(aggr)}').")
+
+        self.aggr_module = aggr_resolver(aggr, **(aggr_kwargs or {}))
 
         self.flow = flow
 
diff --git a/torch_geometric/nn/resolver.py b/torch_geometric/nn/resolver.py
index 099ec3ff7..0bba264f2 100644
--- a/torch_geometric/nn/resolver.py
+++ b/torch_geometric/nn/resolver.py
@@ -63,6 +63,9 @@ def normalization_resolver(query: Union[Any, str], *args, **kwargs):
 
 def aggregation_resolver(query: Union[Any, str], *args, **kwargs):
     import torch_geometric.nn.aggr as aggr
+    if isinstance(query, (list, tuple)):
+        return aggr.MultiAggregation(query, *args, **kwargs)
+
     base_cls = aggr.Aggregation
     aggrs = [
         aggr for aggr in vars(aggr).values()


Test patch:
diff --git a/test/nn/test_resolver.py b/test/nn/test_resolver.py
index cc3dac69e..321df85d4 100644
--- a/test/nn/test_resolver.py
+++ b/test/nn/test_resolver.py
@@ -23,16 +23,16 @@ def test_activation_resolver():
 
 
 @pytest.mark.parametrize('aggr_tuple', [
-    (torch_geometric.nn.aggr.MeanAggregation, 'mean'),
-    (torch_geometric.nn.aggr.SumAggregation, 'sum'),
-    (torch_geometric.nn.aggr.SumAggregation, 'add'),
-    (torch_geometric.nn.aggr.MaxAggregation, 'max'),
-    (torch_geometric.nn.aggr.MinAggregation, 'min'),
-    (torch_geometric.nn.aggr.MulAggregation, 'mul'),
-    (torch_geometric.nn.aggr.VarAggregation, 'var'),
-    (torch_geometric.nn.aggr.StdAggregation, 'std'),
-    (torch_geometric.nn.aggr.SoftmaxAggregation, 'softmax'),
-    (torch_geometric.nn.aggr.PowerMeanAggregation, 'powermean'),
+    (torch_geometric.nn.MeanAggregation, 'mean'),
+    (torch_geometric.nn.SumAggregation, 'sum'),
+    (torch_geometric.nn.SumAggregation, 'add'),
+    (torch_geometric.nn.MaxAggregation, 'max'),
+    (torch_geometric.nn.MinAggregation, 'min'),
+    (torch_geometric.nn.MulAggregation, 'mul'),
+    (torch_geometric.nn.VarAggregation, 'var'),
+    (torch_geometric.nn.StdAggregation, 'std'),
+    (torch_geometric.nn.SoftmaxAggregation, 'softmax'),
+    (torch_geometric.nn.PowerMeanAggregation, 'powermean'),
 ])
 def test_aggregation_resolver(aggr_tuple):
     aggr_module, aggr_repr = aggr_tuple
@@ -40,16 +40,28 @@ def test_aggregation_resolver(aggr_tuple):
     assert isinstance(aggregation_resolver(aggr_repr), aggr_module)
 
 
+def test_multi_aggregation_resolver():
+    aggr = aggregation_resolver(None)
+    assert aggr is None
+
+    aggr = aggregation_resolver(['sum', 'mean', None])
+    assert isinstance(aggr, torch_geometric.nn.MultiAggregation)
+    assert len(aggr.aggrs) == 3
+    assert isinstance(aggr.aggrs[0], torch_geometric.nn.SumAggregation)
+    assert isinstance(aggr.aggrs[1], torch_geometric.nn.MeanAggregation)
+    assert aggr.aggrs[2] is None
+
+
 @pytest.mark.parametrize('norm_tuple', [
-    (torch_geometric.nn.norm.BatchNorm, 'batch', (16, )),
-    (torch_geometric.nn.norm.BatchNorm, 'batch_norm', (16, )),
-    (torch_geometric.nn.norm.InstanceNorm, 'instance_norm', (16, )),
-    (torch_geometric.nn.norm.LayerNorm, 'layer_norm', (16, )),
-    (torch_geometric.nn.norm.GraphNorm, 'graph_norm', (16, )),
-    (torch_geometric.nn.norm.GraphSizeNorm, 'graphsize_norm', ()),
-    (torch_geometric.nn.norm.PairNorm, 'pair_norm', ()),
-    (torch_geometric.nn.norm.MessageNorm, 'message_norm', ()),
-    (torch_geometric.nn.norm.DiffGroupNorm, 'diffgroup_norm', (16, 4)),
+    (torch_geometric.nn.BatchNorm, 'batch', (16, )),
+    (torch_geometric.nn.BatchNorm, 'batch_norm', (16, )),
+    (torch_geometric.nn.InstanceNorm, 'instance_norm', (16, )),
+    (torch_geometric.nn.LayerNorm, 'layer_norm', (16, )),
+    (torch_geometric.nn.GraphNorm, 'graph_norm', (16, )),
+    (torch_geometric.nn.GraphSizeNorm, 'graphsize_norm', ()),
+    (torch_geometric.nn.PairNorm, 'pair_norm', ()),
+    (torch_geometric.nn.MessageNorm, 'message_norm', ()),
+    (torch_geometric.nn.DiffGroupNorm, 'diffgroup_norm', (16, 4)),
 ])
 def test_normalization_resolver(norm_tuple):
     norm_module, norm_repr, norm_args = norm_tuple