Add usage of disable_dynamic_shape decorator in aggregation layers

CHANGELOG.md

@@ -10,7 +10,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Enabled `LinkNeighborLoader` to return number of sampled nodes and edges per hop ([#7516](https://github.com/pyg-team/pytorch_geometric/pull/7516))
 - Added the `HM` personalized fashion recommendation dataset ([#7515](https://github.com/pyg-team/pytorch_geometric/pull/7515))
 - Added the `GraphMixer` model ([#7501](https://github.com/pyg-team/pytorch_geometric/pull/7501))
-- Added the `disable_dynamic_shape` experimental flag ([#7246](https://github.com/pyg-team/pytorch_geometric/pull/7246))
+- Added the `disable_dynamic_shape` experimental flag ([#7246](https://github.com/pyg-team/pytorch_geometric/pull/7246), [#7534](https://github.com/pyg-team/pytorch_geometric/pull/7534))
 - Added the option to override `use_segmm` selection in `HeteroLinear` ([#7474](https://github.com/pyg-team/pytorch_geometric/pull/7474))
 - Added the `MovieLens-1M` heterogeneous dataset ([#7479](https://github.com/pyg-team/pytorch_geometric/pull/7479))
 - Added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493))

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
 
 
@@ -91,6 +92,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/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,