[ONNX] Support aten::scatter_reduce

[vanly89]

🐛 Describe the bug

Hi, experts,
Thanks for your dedication to the great project pytorch. When I uesd torch.onnx.export() function to try to convert .pt model to .onnx model, I got the an error : " ONNX export failure: Exporting the operator ::scatter_reduce to ONNX opset version 16 is not supported. ". Yes, I used torch.scatter_reduce ops in my model. So, did torch.onnx.export support converting scatter_reduce operator to onnx format, if not, is there any other way to replace scatter_reduce operator that torch.onnx export function supported ? Could you so nice to give some guidiance？
I simplified my code as follows：

import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, x, index, input):
        y_max = input.scatter_reduce(0, index, x, reduce='amax')
        y_sum = input.scatter_reduce(0, index, x, reduce='sum')
        return y_max, y_sum


model = Model()
model.eval()

src = torch.tensor([1., 2., 3., 4., 5., 6.])
index = torch.tensor([0, 1, 0, 1, 2, 1])
input = torch.tensor([1., 2., 3., 8.])
torch.onnx.export(model,
                  (src, index, input),
                  'model.onnx',
                  opset_version=16
                  )

Versions

[pip3] numpy==1.19.0
[pip3] pytorch-lightning==1.7.2
[pip3] torch==1.13.0.dev20220828+cu113
[pip3] torch-cluster==1.6.0
[pip3] torch-geometric==2.1.0
[pip3] torch-scatter==2.0.9
[pip3] torch-sparse==0.6.15
[pip3] torch-spline-conv==1.2.1
[pip3] torchaudio==0.12.1
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.13.1
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.3.1 h2bc3f7f_2
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py38h7f8727e_0
[conda] mkl_fft 1.3.1 py38hd3c417c_0
[conda] mkl_random 1.2.2 py38h51133e4_0
[conda] numpy 1.19.0 pypi_0 pypi
[conda] pytorch-lightning 1.7.2 pypi_0 pypi
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torch 1.13.0.dev20220828+cu113 pypi_0 pypi
[conda] torch-cluster 1.6.0 pypi_0 pypi
[conda] torch-geometric 2.1.0 pypi_0 pypi
[conda] torch-scatter 2.0.9 pypi_0 pypi
[conda] torch-sparse 0.6.15 pypi_0 pypi
[conda] torch-spline-conv 1.2.1 pypi_0 pypi
[conda] torchaudio 0.12.1 py38_cu113 pytorch
[conda] torchmetrics 0.9.3 pypi_0 pypi
[conda] torchvision 0.13.1 py38_cu113 pytorch

[justinchuby]

Thanks for your question. scatter_reduce is not yet supported by torch.onnx. We have added it to our list. #84496

We welcome your contribution!

[ThomasRetornaz]

Thanks for this great project
Important note this limitation continue to exist in torch2.0/opset18 and block pygeometric export eg of a simple GATConv
Which is very annoying

[LTsommer]

Thanks for this great project Important note this limitation continue to exist in torch2.0/opset18 and block pygeometric export eg of a simple GATConv Which is very annoying

Same problem, so have you solved it?