Add usage disable_dynamic_shape decorator in aggregation layers.

CHANGELOG.md

@@ -94,6 +94,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added an optional `batch_size` and `max_num_nodes` arguments to `MemPooling` layer ([#7239](https://github.com/pyg-team/pytorch_geometric/pull/7239))
 - Fixed training issues of the GraphGPS example ([#7377](https://github.com/pyg-team/pytorch_geometric/pull/7377))
 - Allowed `CaptumExplainer` to be called multiple times in a row ([#7391](https://github.com/pyg-team/pytorch_geometric/pull/7391))
+- Added usage of disable_dynamic_shape decorator in aggregation layers ([#7534](https://github.com/pyg-team/pytorch_geometric/pull/7534))
+
 
 ### Removed
 

torch_geometric/nn/aggr/base.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.utils import scatter, segment, to_dense_batch
 
 
@@ -58,6 +59,7 @@ class Aggregation(torch.nn.Module):
         - **output:** graph features :math:`(*, |\mathcal{G}|, F_{out})` or
           node features :math:`(*, |\mathcal{V}|, F_{out})`
     """
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,
@@ -91,6 +93,7 @@ def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         pass
 
+    @disable_dynamic_shapes(required_args=['dim_size'])
     def __call__(self, x: Tensor, index: Optional[Tensor] = None,
                  ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                  dim: int = -2, **kwargs) -> Tensor:

torch_geometric/nn/aggr/gmt.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.aggr.utils import (
     PoolingByMultiheadAttention,
@@ -65,6 +66,7 @@ def reset_parameters(self):
             encoder.reset_parameters()
         self.pma2.reset_parameters()
 
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,

torch_geometric/nn/aggr/gru.py

@@ -3,6 +3,7 @@
 from torch import Tensor
 from torch.nn import GRU
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 
 
@@ -29,6 +30,7 @@ def __init__(self, in_channels: int, out_channels: int, **kwargs):
     def reset_parameters(self):
         self.gru.reset_parameters()
 
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,

torch_geometric/nn/aggr/lstm.py

@@ -3,6 +3,7 @@
 from torch import Tensor
 from torch.nn import LSTM
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 
 
@@ -29,6 +30,7 @@ def __init__(self, in_channels: int, out_channels: int, **kwargs):
     def reset_parameters(self):
         self.lstm.reset_parameters()
 
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,

torch_geometric/nn/aggr/mlp.py

@@ -2,6 +2,7 @@
 
 from torch import Tensor
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 
 
@@ -22,6 +23,7 @@ class MLPAggregation(Aggregation):
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.models.MLP`.
     """
+    @disable_dynamic_shapes(required_args=['max_num_elements'])
     def __init__(self, in_channels: int, out_channels: int,
                  max_num_elements: int, **kwargs):
         super().__init__()

torch_geometric/nn/aggr/set_transformer.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.aggr.utils import (
     PoolingByMultiheadAttention,
@@ -73,6 +74,7 @@ def reset_parameters(self):
         for decoder in self.decoders:
             decoder.reset_parameters()
 
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,

torch_geometric/nn/aggr/sort.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+from torch_geometric.experimental import disable_dynamic_shapes
 from torch_geometric.nn.aggr import Aggregation
 
 
@@ -20,6 +21,7 @@ def __init__(self, k: int):
         super().__init__()
         self.k = k
 
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
     def forward(
         self,
         x: Tensor,