Repository: pyg-team/pytorch_geometric
Pull request #7077 (2023-03-29 20:16:15)

`MultiAggr` w/ `HeteroDictLinear`

rusty1s: Related PRs:
* https://github.com/pytorch/pytorch/pull/97960
* https://github.com/pyg-team/pyg-lib/pull/220

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 65f00b012..3894f5f54 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Accelerated attention-based `MultiAggregation` ([#7077](https://github.com/pyg-team/pytorch_geometric/pull/7077))
 - Edges in `HeterophilousGraphDataset` are now undirected by default ([#7065](https://github.com/pyg-team/pytorch_geometric/pull/7065))
 - Fixed a bug in `FastHGTConv` that computed values via parameters used to compute the keys ([#7050](https://github.com/pyg-team/pytorch_geometric/pull/7050))
 - Accelerated sparse tensor conversion routines ([#7042](https://github.com/pyg-team/pytorch_geometric/pull/7042), [#7043](https://github.com/pyg-team/pytorch_geometric/pull/7043))
diff --git a/torch_geometric/nn/aggr/multi.py b/torch_geometric/nn/aggr/multi.py
index 6883f935d..72b941ced 100644
--- a/torch_geometric/nn/aggr/multi.py
+++ b/torch_geometric/nn/aggr/multi.py
@@ -7,6 +7,7 @@ from torch.nn import Linear, MultiheadAttention
 
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.aggr.fused import FusedAggregation
+from torch_geometric.nn.dense import HeteroDictLinear
 from torch_geometric.nn.resolver import aggregation_resolver
 
 
@@ -114,10 +115,8 @@ class MultiAggregation(Aggregation):
                 )
 
             elif mode == 'attn':
-                self.lin_heads = torch.nn.ModuleList([
-                    Linear(channels, self.out_channels)
-                    for channels in self.in_channels
-                ])
+                channels = {str(k): v for k, v, in enumerate(self.in_channels)}
+                self.lin_heads = HeteroDictLinear(channels, self.out_channels)
                 num_heads = mode_kwargs.pop('num_heads', 1)
                 self.multihead_attn = MultiheadAttention(
                     self.out_channels,
@@ -137,8 +136,7 @@ class MultiAggregation(Aggregation):
         if self.mode == 'proj':
             self.lin.reset_parameters()
         if self.mode == 'attn':
-            for lin in self.lin_heads:
-                lin.reset_parameters()
+            self.lin_heads.reset_parameters()
             self.multihead_attn._reset_parameters()
 
     def get_out_channels(self, in_channels: int) -> int:
@@ -181,10 +179,10 @@ class MultiAggregation(Aggregation):
             return self.lin(torch.cat(inputs, dim=-1))
 
         if hasattr(self, 'multihead_attn'):
-            x = torch.stack(
-                [head(inputs[i]) for i, head in enumerate(self.lin_heads)],
-                dim=0,
-            )
+            x_dict = {str(k): v for k, v, in enumerate(inputs)}
+            x_dict = self.lin_heads(x_dict)
+            xs = [x_dict[str(key)] for key in range(len(inputs))]
+            x = torch.stack(xs, dim=0)
             attn_out, _ = self.multihead_attn(x, x, x)
             return torch.mean(attn_out, dim=0)
 
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index c8971e1c9..4567bc909 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -369,25 +369,22 @@ class HeteroDictLinear(torch.nn.Module):
 
     def forward(
         self,
-        x_dict: Dict[Any, Tensor],
-    ) -> Dict[Any, Tensor]:
+        x_dict: Dict[str, Tensor],
+    ) -> Dict[str, Tensor]:
         r"""
         Args:
             x_dict (Dict[Any, torch.Tensor]): A dictionary holding input
                 features for each individual type.
         """
         if torch_geometric.typing.WITH_GMM:
-            xs = [x_dict[key] for key in x_dict.keys()]
-            weights = [self.lins[key].weight.t() for key in x_dict.keys()]
-            if self.kwargs.get('bias', True):
-                biases = [self.lins[key].bias for key in x_dict.keys()]
-            else:
-                biases = None
-
+            xs = [x_dict[key] for key in self.lins.keys()]
+            weights = [lin.weight.t() for lin in self.lins.values()]
+            biases = [lin.bias for lin in self.lins.values()]
+            biases = None if biases[0] is None else biases
             outs = pyg_lib.ops.grouped_matmul(xs, weights, biases)
             return {key: out for key, out in zip(x_dict.keys(), outs)}
 
-        return {key: self.lins[key](x) for key, x in x_dict.items()}
+        return {key: lin(x_dict[key]) for key, lin in self.lins.items()}
 
     @torch.no_grad()
     def initialize_parameters(self, module, input):


Test patch:
diff --git a/test/nn/aggr/test_multi.py b/test/nn/aggr/test_multi.py
index 222696b62..e37eb3923 100644
--- a/test/nn/aggr/test_multi.py
+++ b/test/nn/aggr/test_multi.py
@@ -45,5 +45,10 @@ def test_multi_aggr(multi_aggr_tuple):
     else:
         assert torch.allclose(out, aggr(x, ptr=ptr))
 
-    jit = torch.jit.script(aggr)
-    assert torch.allclose(out, jit(x, index))
+    if aggr_kwargs['mode'] == 'attn' and torch_geometric.typing.WITH_GMM:
+        # See: https://github.com/pytorch/pytorch/pull/97960
+        with pytest.raises(RuntimeError, match="Unknown builtin op"):
+            jit = torch.jit.script(aggr)
+    else:
+        jit = torch.jit.script(aggr)
+        assert torch.allclose(out, jit(x, index))
diff --git a/test/nn/conv/test_graph_conv.py b/test/nn/conv/test_graph_conv.py
index 14bd30481..082e2954f 100644
--- a/test/nn/conv/test_graph_conv.py
+++ b/test/nn/conv/test_graph_conv.py
@@ -47,8 +47,8 @@ def test_graph_conv():
     if is_full_test() and torch_geometric.typing.WITH_TORCH_SPARSE:
         t = '(Tensor, SparseTensor, OptTensor, Size) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert torch.allclose(jit(x1, adj3.t()), out1)
-        assert torch.allclose(jit(x1, adj4.t()), out2)
+        assert torch.allclose(jit(x1, adj3.t()), out1, atol=1e-6)
+        assert torch.allclose(jit(x1, adj4.t()), out2, atol=1e-6)
 
     # Test bipartite message passing:
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
diff --git a/test/nn/conv/test_signed_conv.py b/test/nn/conv/test_signed_conv.py
index 6e72e29f9..8ee73daf0 100644
--- a/test/nn/conv/test_signed_conv.py
+++ b/test/nn/conv/test_signed_conv.py
@@ -50,8 +50,10 @@ def test_signed_conv():
     # Test bipartite message passing:
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
 
-    assert torch.allclose(conv1((x, x[:2]), edge_index, edge_index), out1[:2])
-    assert torch.allclose(conv1((x, x[:2]), adj1.t(), adj1.t()), out1[:2])
+    assert torch.allclose(conv1((x, x[:2]), edge_index, edge_index), out1[:2],
+                          atol=1e-6)
+    assert torch.allclose(conv1((x, x[:2]), adj1.t(), adj1.t()), out1[:2],
+                          atol=1e-6)
     assert torch.allclose(conv2((out1, out1[:2]), edge_index, edge_index),
                           out2[:2], atol=1e-6)
     assert torch.allclose(conv2((out1, out1[:2]), adj1.t(), adj1.t()),
@@ -59,7 +61,8 @@ def test_signed_conv():
 
     if torch_geometric.typing.WITH_TORCH_SPARSE:
         adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 2))
-        assert torch.allclose(conv1((x, x[:2]), adj2.t(), adj2.t()), out1[:2])
+        assert torch.allclose(conv1((x, x[:2]), adj2.t(), adj2.t()), out1[:2],
+                              atol=1e-6)
         assert torch.allclose(conv2((out1, out1[:2]), adj2.t(), adj2.t()),
                               out2[:2], atol=1e-6)
 
@@ -76,6 +79,7 @@ def test_signed_conv():
         t = '(PairTensor, SparseTensor, SparseTensor) -> Tensor'
         jit1 = torch.jit.script(conv1.jittable(t))
         jit2 = torch.jit.script(conv2.jittable(t))
-        assert torch.allclose(jit1((x, x[:2]), adj2.t(), adj2.t()), out1[:2])
+        assert torch.allclose(jit1((x, x[:2]), adj2.t(), adj2.t()), out1[:2],
+                              atol=1e-6)
         assert torch.allclose(jit2((out1, out1[:2]), adj2.t(), adj2.t()),
                               out2[:2], atol=1e-6)
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index c74c13d27..6450d198e 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -5,8 +5,9 @@ import torch
 from torch.nn import Linear as PTLinear
 from torch.nn.parameter import UninitializedParameter
 
+import torch_geometric.typing
 from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
-from torch_geometric.testing import is_full_test, withPackage
+from torch_geometric.testing import withPackage
 
 weight_inits = ['glorot', 'kaiming_uniform', None]
 bias_inits = ['zeros', None]
@@ -111,9 +112,8 @@ def test_hetero_linear():
     out = lin(x, type_vec)
     assert out.size() == (3, 32)
 
-    if is_full_test():
-        jit = torch.jit.script(lin)
-        assert torch.allclose(jit(x, type_vec), out)
+    jit = torch.jit.script(lin)
+    assert torch.allclose(jit(x, type_vec), out)
 
 
 def test_lazy_hetero_linear():
@@ -127,11 +127,13 @@ def test_lazy_hetero_linear():
     assert out.size() == (3, 32)
 
 
-def test_hetero_dict_linear():
+@pytest.mark.parametrize('bias', [True, False])
+def test_hetero_dict_linear(bias):
     x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 8)}
 
-    lin = HeteroDictLinear({'v': 16, 'w': 8}, 32)
-    assert str(lin) == "HeteroDictLinear({'v': 16, 'w': 8}, 32, bias=True)"
+    lin = HeteroDictLinear({'v': 16, 'w': 8}, 32, bias=bias)
+    assert str(lin) == (f"HeteroDictLinear({{'v': 16, 'w': 8}}, 32, "
+                        f"bias={bias})")
 
     out_dict = lin(x_dict)
     assert len(out_dict) == 2
@@ -140,14 +142,23 @@ def test_hetero_dict_linear():
 
     x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 16)}
 
-    lin = HeteroDictLinear(16, 32, types=['v', 'w'])
-    assert str(lin) == "HeteroDictLinear({'v': 16, 'w': 16}, 32, bias=True)"
+    lin = HeteroDictLinear(16, 32, types=['v', 'w'], bias=bias)
+    assert str(lin) == (f"HeteroDictLinear({{'v': 16, 'w': 16}}, 32, "
+                        f"bias={bias})")
 
     out_dict = lin(x_dict)
     assert len(out_dict) == 2
     assert out_dict['v'].size() == (3, 32)
     assert out_dict['w'].size() == (2, 32)
 
+    if torch_geometric.typing.WITH_GMM:
+        # See: https://github.com/pytorch/pytorch/pull/97960
+        with pytest.raises(RuntimeError, match="Unknown builtin op"):
+            jit = torch.jit.script(lin)
+    else:
+        jit = torch.jit.script(lin)
+        assert len(jit(x_dict)) == 2
+
 
 def test_lazy_hetero_dict_linear():
     x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 8)}