Repository: pyg-team/pytorch_geometric

Issue:
TorchScript attentional aggregation error
I'm trying to compile gnn model with torch script and get following error:

```bash
RuntimeError: 
Unknown type name 'torch_geometric.nn.MLP':
  File "/usr/local/lib/python3.10/dist-packages/torch_geometric/nn/aggr/attention.py", line 55
        self.assert_two_dimensional_input(x, dim)
    
        if isinstance(self.gate_nn, torch_geometric.nn.MLP):
                                    ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            gate = self.gate_nn(x, index, dim_size)
        else:
```

Here is example to reproduce this error:

```python
import torch
from torch import nn
from torch_geometric.nn import aggr, TransformerConv
from torch_geometric.nn.norm import GraphNorm


class GNN(nn.Module):
    def __init__(self):
        super(GNN, self).__init__()

        self.gnn = TransformerConv(10, 10, 1, dropout=0.1).jittable()         
        self.normalization = GraphNorm(in_channels=10)
        self.act = nn.GELU()
        self.aggr = aggr.AttentionalAggregation(gate_nn=nn.Linear(10, 1), nn=nn.Linear(10, 10)) #doesn't compile

    def forward(self, x, edge_index, batch_ptr):
        hidden = x
        hidden = self.gnn(hidden, edge_index)
        hidden = self.normalization(hidden)
        hidden = self.act(hidden)

        graph_rep = self.aggr(hidden, ptr=batch_ptr)

        return graph_rep

    
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = GNN().to(device)

x = torch.randn(100, 10).to(device)
edge_index = torch.randint(100, size=(2, 20)).to(device)
batch_ptr = torch.tensor([0, 100]).to(device)
res = model(x, edge_index, batch_ptr)
print(res)

model = torch.jit.script(model)
```

### Environment

* PyG version: 2.4.0
* PyTorch version: 2.1.0
* OS: ubuntu 22
* Python version: 3.10
* CUDA/cuDNN version: 12.3
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip/source


Discussion:
also need this

Resolving patch:
diff --git a/torch_geometric/nn/aggr/attention.py b/torch_geometric/nn/aggr/attention.py
--- a/torch_geometric/nn/aggr/attention.py
+++ b/torch_geometric/nn/aggr/attention.py
@@ -3,7 +3,6 @@
 import torch
 from torch import Tensor
 
-import torch_geometric
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.inits import reset
 from torch_geometric.utils import softmax
@@ -35,16 +34,32 @@ class AttentionalAggregation(Aggregation):
             before combining them with the attention scores, *e.g.*, defined by
             :class:`torch.nn.Sequential`. (default: :obj:`None`)
     """
-    def __init__(self, gate_nn: torch.nn.Module,
-                 nn: Optional[torch.nn.Module] = None):
+    def __init__(
+        self,
+        gate_nn: torch.nn.Module,
+        nn: Optional[torch.nn.Module] = None,
+    ):
         super().__init__()
-        self.gate_nn = gate_nn
-        self.nn = nn
-        self.reset_parameters()
+
+        from torch_geometric.nn import MLP
+
+        self.gate_nn = self.gate_mlp = None
+        if isinstance(gate_nn, MLP):
+            self.gate_mlp = gate_nn
+        else:
+            self.gate_nn = gate_nn
+
+        self.nn = self.mlp = None
+        if isinstance(nn, MLP):
+            self.mlp = nn
+        else:
+            self.nn = nn
 
     def reset_parameters(self):
         reset(self.gate_nn)
+        reset(self.gate_mlp)
         reset(self.nn)
+        reset(self.mlp)
 
     def forward(self, x: Tensor, index: Optional[Tensor] = None,
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
@@ -52,13 +67,13 @@ def forward(self, x: Tensor, index: Optional[Tensor] = None,
 
         self.assert_two_dimensional_input(x, dim)
 
-        if isinstance(self.gate_nn, torch_geometric.nn.MLP):
-            gate = self.gate_nn(x, index, dim_size)
+        if self.gate_mlp is not None:
+            gate = self.gate_mlp(x, batch=index, batch_size=dim_size)
         else:
             gate = self.gate_nn(x)
 
-        if isinstance(self.nn, torch_geometric.nn.MLP):
-            x = self.nn(x, index, dim_size)
+        if self.mlp is not None:
+            x = self.mlp(x, batch=index, batch_size=dim_size)
         elif self.nn is not None:
             x = self.nn(x)
 
@@ -66,5 +81,6 @@ def forward(self, x: Tensor, index: Optional[Tensor] = None,
         return self.reduce(gate * x, index, ptr, dim_size, dim)
 
     def __repr__(self) -> str:
-        return (f'{self.__class__.__name__}(gate_nn={self.gate_nn}, '
-                f'nn={self.nn})')
+        return (f'{self.__class__.__name__}('
+                f'gate_nn={self.gate_mlp or self.gate_nn}, '
+                f'nn={self.mlp or self.nn})')
diff --git a/torch_geometric/nn/aggr/deep_sets.py b/torch_geometric/nn/aggr/deep_sets.py
--- a/torch_geometric/nn/aggr/deep_sets.py
+++ b/torch_geometric/nn/aggr/deep_sets.py
@@ -30,25 +30,46 @@ def __init__(
         global_nn: Optional[torch.nn.Module] = None,
     ):
         super().__init__()
-        self.local_nn = local_nn
-        self.global_nn = global_nn
+
+        from torch_geometric.nn import MLP
+
+        self.local_nn = self.local_mlp = None
+        if isinstance(local_nn, MLP):
+            self.local_mlp = local_nn
+        else:
+            self.local_nn = local_nn
+
+        self.global_nn = self.global_mlp = None
+        if isinstance(global_nn, MLP):
+            self.global_mlp = global_nn
+        else:
+            self.global_nn = global_nn
 
     def reset_parameters(self):
-        if self.local_nn is not None:
-            reset(self.local_nn)
-        if self.global_nn is not None:
-            reset(self.global_nn)
+        reset(self.local_nn)
+        reset(self.local_mlp)
+        reset(self.global_nn)
+        reset(self.global_mlp)
 
     def forward(self, x: Tensor, index: Optional[Tensor] = None,
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
+
+        if self.local_mlp is not None:
+            x = self.local_mlp(x, batch=index, batch_size=dim_size)
         if self.local_nn is not None:
             x = self.local_nn(x)
+
         x = self.reduce(x, index, ptr, dim_size, dim, reduce='sum')
-        if self.global_nn is not None:
+
+        if self.global_mlp is not None:
+            x = self.global_mlp(x, batch=index, batch_size=dim_size)
+        elif self.global_nn is not None:
             x = self.global_nn(x)
+
         return x
 
     def __repr__(self) -> str:
-        return (f'{self.__class__.__name__}(local_nn={self.local_nn}, '
-                f'global_nn={self.global_nn})')
+        return (f'{self.__class__.__name__}('
+                f'local_nn={self.local_mlp or self.local_nn}, '
+                f'global_nn={self.global_mlp or self.global_nn})')
diff --git a/torch_geometric/nn/aggr/mlp.py b/torch_geometric/nn/aggr/mlp.py
--- a/torch_geometric/nn/aggr/mlp.py
+++ b/torch_geometric/nn/aggr/mlp.py
@@ -32,8 +32,13 @@ class MLPAggregation(Aggregation):
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.models.MLP`.
     """
-    def __init__(self, in_channels: int, out_channels: int,
-                 max_num_elements: int, **kwargs):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        max_num_elements: int,
+        **kwargs,
+    ):
         super().__init__()
 
         self.in_channels = in_channels
@@ -41,8 +46,11 @@ def __init__(self, in_channels: int, out_channels: int,
         self.max_num_elements = max_num_elements
 
         from torch_geometric.nn import MLP
-        self.mlp = MLP(in_channels=in_channels * max_num_elements,
-                       out_channels=out_channels, **kwargs)
+        self.mlp = MLP(
+            in_channels=in_channels * max_num_elements,
+            out_channels=out_channels,
+            **kwargs,
+        )
 
         self.reset_parameters()
 
@@ -52,8 +60,10 @@ def reset_parameters(self):
     def forward(self, x: Tensor, index: Optional[Tensor] = None,
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
+
         x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim,
                                    max_num_elements=self.max_num_elements)
+
         return self.mlp(x.view(-1, x.size(1) * x.size(2)), index, dim_size)
 
     def __repr__(self) -> str:


Test patch:
diff --git a/test/nn/aggr/test_attention.py b/test/nn/aggr/test_attention.py
--- a/test/nn/aggr/test_attention.py
+++ b/test/nn/aggr/test_attention.py
@@ -15,6 +15,7 @@ def test_attentional_aggregation():
     gate_nn = MLP([channels, 1], act='relu')
     nn = MLP([channels, channels], act='relu')
     aggr = AttentionalAggregation(gate_nn, nn)
+    aggr.reset_parameters()
     assert str(aggr) == (f'AttentionalAggregation(gate_nn=MLP({channels}, 1), '
                          f'nn=MLP({channels}, {channels}))')