Repository: pyg-team/pytorch_geometric

Issue:
ASAPooling: jit.save function model export failed
### 🐛 Describe the bug

ASAPooling: jit.save function model export failed. I need your help, thank you very much!

```python
class ASAP(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.asapooling = ASAPooling(128, ratio=0.25, dropout=0.0).jittable()

    def forward(self, x, edge_index):

        edge_weight = None
        batch = torch.zeros((x.shape[0])).cuda().long()
        x, edge_index, edge_weight, batch, _ = self.asapooling(
            x=x, edge_index=edge_index, edge_weight=edge_weight,
            batch=batch)

        return x
model = ASAP()
mm = torch.jit.script(model)
torch.jit.save(mm,"asapooling.pt")
```
#error 
Traceback (most recent call last):
  File "D:\pycharm\PyCharm Community Edition 2020.3.5\plugins\python-ce\helpers\pydev\pydevd.py", line 1477, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "D:\pycharm\PyCharm Community Edition 2020.3.5\plugins\python-ce\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "E:/python_program/mesh_gcn/teeth_classfication/net/jit_model.py", line 210, in <module>
    torch.jit.save(mm,"asapooling.pt")
  File "D:\Anaconda3\envs\teeth_deploy\lib\site-packages\torch\jit\_serialization.py", line 81, in save
    m.save(f, _extra_files=_extra_files)
  File "D:\Anaconda3\envs\teeth_deploy\lib\site-packages\torch\jit\_script.py", line 693, in save
    return self._c.save(str(f), **kwargs)
RuntimeError: 
Could not export Python function call '<python_value>'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
  File "E:\python_program\mesh_gcn\teeth_classfication\torch_geometric\utils\scatter.py", line 66
            out: Optional[Tensor] = None, dim_size: Optional[int] = None,
            reduce: str = 'sum') -> Tensor:
    return _scatter_impl(src, index, dim, out, dim_size, reduce)
           ~~~~~~~~~~~~~ <--- HERE


Process finished with exit code -1




### Environment

PyG version:2.04
PyTorch version:1.11.0
OS: windows
Python version:3.7.0
CUDA/cuDNN version: cuda=11.3.1 cudnn not have
How you installed PyTorch and PyG (conda, pip, source): conda
Any other relevant information (e.g., version of torch-scatter):
pytorch-scatter==2.0.9
pytorch-sparse==0.6.13
pytorch-cluster==1.6.0



Discussion:

Resolving patch:
diff --git a/benchmark/inference/inference_benchmark.py b/benchmark/inference/inference_benchmark.py
--- a/benchmark/inference/inference_benchmark.py
+++ b/benchmark/inference/inference_benchmark.py
@@ -3,7 +3,6 @@
 import torch
 from utils import get_dataset, get_model
 
-from torch_geometric import set_experimental_mode
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import PNAConv
 from torch_geometric.profile import rename_profile_file, timeit, torch_profile
@@ -99,13 +98,12 @@ def run(args: argparse.ArgumentParser) -> None:
                         model.eval()
 
                         with amp:
-                            with set_experimental_mode(args.experimental_mode):
-                                for _ in range(args.warmup):
-                                    model.inference(subgraph_loader, device,
-                                                    progress_bar=True)
-                                with timeit():
-                                    model.inference(subgraph_loader, device,
-                                                    progress_bar=True)
+                            for _ in range(args.warmup):
+                                model.inference(subgraph_loader, device,
+                                                progress_bar=True)
+                            with timeit():
+                                model.inference(subgraph_loader, device,
+                                                progress_bar=True)
 
                             if args.profile:
                                 with torch_profile():
@@ -143,8 +141,6 @@ def run(args: argparse.ArgumentParser) -> None:
         '--hetero-num-neighbors', default=10, type=int,
         help='number of neighbors to sample per layer for hetero workloads')
     argparser.add_argument('--num-workers', default=0, type=int)
-    argparser.add_argument('--experimental-mode', action='store_true',
-                           help='use experimental mode')
     argparser.add_argument('--warmup', default=1, type=int)
     argparser.add_argument('--profile', action='store_true')
     argparser.add_argument('--bf16', action='store_true')
diff --git a/torch_geometric/experimental.py b/torch_geometric/experimental.py
--- a/torch_geometric/experimental.py
+++ b/torch_geometric/experimental.py
@@ -1,6 +1,6 @@
 from typing import List, Optional, Union
 
-__experimental_flag__ = {'scatter_reduce': False}
+__experimental_flag__ = {}
 
 Options = Optional[Union[str, List[str]]]
 
@@ -36,11 +36,8 @@ class experimental_mode:
         ...     out = model(data.x, data.edge_index)
 
     Args:
-        options (str or list, optional): Possible option(s):
-
-            - :obj:`"torch_scatter"`: Enables the usage of
-              :meth:`torch.scatter_reduce` instead of
-              :meth:`torch_scatter.scatter`. Requires :obj:`torch>=1.12`.
+        options (str or list, optional): Currently there are no experimental
+            features.
     """
     def __init__(self, options: Options = None):
         self.options = get_options(options)
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -1,66 +1,94 @@
-import warnings
 from typing import Optional
 
 import torch
-import torch_scatter
 from torch import Tensor
 
-import torch_geometric
-
 major, minor, _ = torch.__version__.split('.', maxsplit=2)
 major, minor = int(major), int(minor)
 has_pytorch112 = major > 1 or (major == 1 and minor >= 12)
 
+if has_pytorch112:  # pragma: no cover
 
-class ScatterImpl(torch.nn.Module):
-    def forward(self, src: Tensor, index: Tensor, dim: int = -1,
+    class ScatterHelpers:
+        @staticmethod
+        def broadcast(src: Tensor, other: Tensor, dim: int) -> Tensor:
+            dim = other.dim() + dim if dim < 0 else dim
+            if src.dim() == 1:
+                for _ in range(0, dim):
+                    src = src.unsqueeze(0)
+            for _ in range(src.dim(), other.dim()):
+                src = src.unsqueeze(-1)
+            src = src.expand(other.size())
+            return src
+
+        @staticmethod
+        def generate_out(src: Tensor, index: Tensor, dim: int,
+                         dim_size: Optional[int]) -> Tensor:
+            size = list(src.size())
+            if dim_size is not None:
+                size[dim] = dim_size
+            elif index.numel() > 0:
+                size[dim] = int(index.max()) + 1
+            else:
+                size[dim] = 0
+            return src.new_zeros(size)
+
+        @staticmethod
+        def scatter_mean(src: Tensor, index: Tensor, dim: int, out: Tensor,
+                         dim_size: Optional[int]) -> Tensor:
+            out.scatter_add_(dim, index, src)
+
+            index_dim = dim
+            if index_dim < 0:
+                # `index_dim` counts axes from the begining (0)
+                index_dim = index_dim + src.dim()
+            if index.dim() <= index_dim:
+                # in case `index` was broadcasted, `count` scatter should be
+                # performed over the last axis
+                index_dim = index.dim() - 1
+
+            ones = torch.ones(index.size(), dtype=src.dtype, device=src.device)
+            count = ScatterHelpers.generate_out(ones, index, index_dim,
+                                                dim_size)
+            count.scatter_add_(index_dim, index, ones)
+            count[count < 1] = 1
+            count = ScatterHelpers.broadcast(count, out, dim)
+            if out.is_floating_point():
+                out.true_divide_(count)
+            else:
+                out.div_(count, rounding_mode='floor')
+            return out
+
+    def scatter(src: Tensor, index: Tensor, dim: int = -1,
+                out: Optional[Tensor] = None, dim_size: Optional[int] = None,
+                reduce: str = 'sum') -> Tensor:
+        reduce = 'sum' if reduce == 'add' else reduce
+        reduce = 'prod' if reduce == 'mul' else reduce
+        reduce = 'amin' if reduce == 'min' else reduce
+        reduce = 'amax' if reduce == 'max' else reduce
+
+        index = ScatterHelpers.broadcast(index, src, dim)
+        include_self = out is not None
+
+        if out is None:  # Generate `out` if not given:
+            out = ScatterHelpers.generate_out(src, index, dim, dim_size)
+
+        # explicit usage of `torch.scatter_add_` and switching to
+        # `torch_scatter` implementation of mean algorithm comes with
+        # significant performance boost.
+        # TODO: use only `torch.scatter_reduce_` after performance issue will
+        # be fixed on the PyTorch side.
+        if reduce == 'mean':
+            return ScatterHelpers.scatter_mean(src, index, dim, out, dim_size)
+        elif reduce == 'sum':
+            return out.scatter_add_(dim, index, src)
+        return out.scatter_reduce_(dim, index, src, reduce,
+                                   include_self=include_self)
+
+else:
+    import torch_scatter
+
+    def scatter(src: Tensor, index: Tensor, dim: int = -1,
                 out: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 reduce: str = 'sum') -> Tensor:
-        if torch_geometric.is_experimental_mode_enabled('scatter_reduce'):
-            if not has_pytorch112:
-                warnings.warn("Cannot use 'scatter_reduce' experimental mode "
-                              "on PyTorch < 1.12")
-                return torch_scatter.scatter(src, index, dim, out, dim_size,
-                                             reduce)
-
-            warnings.filterwarnings('ignore', '.*scatter_reduce.*')
-
-            reduce = 'sum' if reduce == 'add' else reduce
-            reduce = 'prod' if reduce == 'mul' else reduce
-            reduce = 'amin' if reduce == 'min' else reduce
-            reduce = 'amax' if reduce == 'max' else reduce
-
-            # Broadcast `index`:
-            dim = src.dim() + dim if dim < 0 else dim
-            if index.dim() == 1:
-                for _ in range(0, dim):
-                    index = index.unsqueeze(0)
-            for _ in range(index.dim(), src.dim()):
-                index = index.unsqueeze(-1)
-            index = index.expand(src.size())
-
-            include_self = out is not None
-
-            if out is None:  # Generate `out` if not given:
-                size = list(src.size())
-                if dim_size is not None:
-                    size[dim] = dim_size
-                elif index.numel() > 0:
-                    size[dim] = int(index.max()) + 1
-                else:
-                    size[dim] = 0
-                out = src.new_zeros(size)
-
-            return out.scatter_reduce_(dim, index, src, reduce,
-                                       include_self=include_self)
-
         return torch_scatter.scatter(src, index, dim, out, dim_size, reduce)
-
-
-_scatter_impl = ScatterImpl()
-
-
-def scatter(src: Tensor, index: Tensor, dim: int = -1,
-            out: Optional[Tensor] = None, dim_size: Optional[int] = None,
-            reduce: str = 'sum') -> Tensor:
-    return _scatter_impl(src, index, dim, out, dim_size, reduce)


Test patch:
diff --git a/test/nn/pool/test_asap.py b/test/nn/pool/test_asap.py
--- a/test/nn/pool/test_asap.py
+++ b/test/nn/pool/test_asap.py
@@ -1,7 +1,10 @@
+import io
+
+import pytest
 import torch
 
 from torch_geometric.nn import ASAPooling, GCNConv, GraphConv
-from torch_geometric.testing import is_full_test
+from torch_geometric.testing import is_full_test, onlyFullTest
 
 
 def test_asap():
@@ -33,3 +36,14 @@ def test_asap():
         out = pool(x, edge_index)
         assert out[0].size() == (2, in_channels)
         assert out[1].size() == (2, 2)
+
+
+@onlyFullTest
+def test_asap_jit_save():
+    pool = ASAPooling(in_channels=16)
+    pool_jit = pool.jittable()
+    model = torch.jit.script(pool_jit)
+    try:
+        torch.jit.save(model, io.BytesIO())
+    except RuntimeError:
+        pytest.fail('ASAP model serialization failed.')
diff --git a/test/test_experimental.py b/test/test_experimental.py
--- a/test/test_experimental.py
+++ b/test/test_experimental.py
@@ -7,7 +7,8 @@
 )
 
 
-@pytest.mark.parametrize('options', [None, 'scatter_reduce'])
+@pytest.mark.skip(reason='No experimental options available right now.')
+@pytest.mark.parametrize('options', [None])
 def test_experimental_mode(options):
     assert is_experimental_mode_enabled(options) is False
     with experimental_mode(options):
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -2,7 +2,6 @@
 import torch
 import torch_scatter
 
-import torch_geometric
 from torch_geometric.testing import withPackage
 from torch_geometric.utils import scatter
 
@@ -14,8 +13,7 @@ def test_scatter(reduce):
     src = torch.randn(8, 100, 32)
     index = torch.randint(0, 10, (100, ), dtype=torch.long)
 
-    with torch_geometric.experimental_mode('scatter_reduce'):
-        out1 = scatter(src, index, dim=1, reduce=reduce)
+    out1 = scatter(src, index, dim=1, reduce=reduce)
     out2 = torch_scatter.scatter(src, index, dim=1, reduce=reduce)
     assert torch.allclose(out1, out2, atol=1e-6)
 
@@ -27,8 +25,7 @@ def test_pytorch_scatter_backward(reduce):
     src = torch.randn(8, 100, 32).requires_grad_(True)
     index = torch.randint(0, 10, (100, ), dtype=torch.long)
 
-    with torch_geometric.experimental_mode('scatter_reduce'):
-        out = scatter(src, index, dim=1, reduce=reduce).relu()
+    out = scatter(src, index, dim=1, reduce=reduce).relu()
 
     assert src.grad is None
     out.mean().backward()
@@ -36,15 +33,14 @@ def test_pytorch_scatter_backward(reduce):
 
 
 @withPackage('torch>=1.12.0')
-@pytest.mark.parametrize('reduce', ['sum', 'add', 'mean', 'min', 'max'])
+@pytest.mark.parametrize('reduce', ['min', 'max'])
 def test_pytorch_scatter_inplace_backward(reduce):
     torch.manual_seed(12345)
 
     src = torch.randn(8, 100, 32).requires_grad_(True)
     index = torch.randint(0, 10, (100, ), dtype=torch.long)
 
-    with torch_geometric.experimental_mode('scatter_reduce'):
-        out = scatter(src, index, dim=1, reduce=reduce).relu_()
+    out = scatter(src, index, dim=1, reduce=reduce).relu_()
 
     with pytest.raises(RuntimeError, match="modified by an inplace operation"):
         out.mean().backward()
@@ -58,8 +54,7 @@ def test_scatter_with_out(reduce):
     index = torch.randint(0, 10, (100, ), dtype=torch.long)
     out = torch.randn(8, 10, 32)
 
-    with torch_geometric.experimental_mode('scatter_reduce'):
-        out1 = scatter(src, index, dim=1, out=out.clone(), reduce=reduce)
+    out1 = scatter(src, index, dim=1, out=out.clone(), reduce=reduce)
     out2 = torch_scatter.scatter(src, index, dim=1, out=out.clone(),
                                  reduce=reduce)
     assert torch.allclose(out1, out2, atol=1e-6)