scatter max/min

[philass]

This PR proposes adding max and min as supported reduction attributes for both ScatterElements and ScatterND.

Closes #4322

ScatterElements and ScatterND provide a mechanism for specifying the reduction function. Currently only add and mul are supported. It would be useful to extend these to support max and min, as these come up in practice and are supported by both Pytorch and TensorFlow. With out this addition ONNX doesn't have a good way to represent these operations.

[gramalingam]

nit: may be we can combine all these 4 descriptions into one, as:

    output = np.copy(data)
    update_indices = indices.shape[:-1]
    for idx in np.ndindex(update_indices):
        output[indices[idx]] = reduction-op (output[indices[idx]], updates[idx])
where the reduction-op is +/*/max/min as specified.

[philass]

Thats a great suggestion

[gramalingam]

As above: it would help to consolidate all the cases into one.

[gramalingam]

I suggest adding the following, as the last part of the documentation:

(Opset 18 change): Adds max/min to the set of allowed reduction ops.

[p-wysocki]

LGTM, after @gramalingam suggestions. These are all the reduction functions PyTorch uses, but I also found more in Apple's BNNS. We could try figuring out if any of them could be useful and added together with max and min. Just a thought, I believe it's okay as is.

[philass]

@gramalingam, Thanks for the review and the great suggestions. I have updated the PR with your recommendations. Sorry for the delay!

@p-wysocki thats a really good point. There is this delicate balance between trying to not add every possible reduction op and adequately expressing Pytorch and Tensorflow ops in ONNX.

There are more Pytorch Scatter ops than just Add, Mul, Max and Min. We found that the Pytorch scatter package is something that is quite common. It consists of the following Scatters.

• Scatter Add
• Scatter Div
• Scatter Max
• Scatter Mean
• Scatter Min
• Scatter Mul
• Scatter Std
• Scatter Sub

The reason that I only added Max and Min is that we have actually seen models that need ScatterMax and ScatterMin. Namely

• GATConv
• GATv2Conv
• GENConv
• TransformerConv
• SuperGATConv
• AGNNConv

Furthermore Sub, Std, Mean, Div can likely be represented by a combination of ScatterAdd, and ScatterMul, and other ops to get the required behavior.

I think a true general solution is to use some of the ideas of If and Loop where some subgraph can be given as an attribute. This subgraph would be the reduction operation. That is more involved, and I don't have the knowledge of ONNX to drive that at the moment.

In the short term I think Max and Min are a good start.

[p-wysocki]

@philass Thank you for such a detailed answer, I agree that Max and Min are sufficient for now.

[gramalingam]

LGTM, thanks!

[gramalingam]

Hi @philass : thanks again for this PR. It seems to require some merge conflict (hopefully just the minor one in operator_sets.h). Once that is addressed, we can merge the PR.

[philass]

@gramalingam thanks for pointing that out.

I've resolved the conflict, and I am ready to get this merged :)

[PallottaEnrico]

Hi @philass , i'm using the latest pytorch version, however i'm still facing the error:
torch.onnx.errors.UnsupportedOperatorError: ONNX export failed on an operator with unrecognized namespace torch_scatter::scatter_max. If you are trying to export a custom operator, make sure you registered it with the right domain and version.

Shouldn't it be solved?

[LebronRemonJames]

Hi @philass , i'm using the latest pytorch version, however i'm still facing the error: torch.onnx.errors.UnsupportedOperatorError: ONNX export failed on an operator with unrecognized namespace torch_scatter::scatter_max. If you are trying to export a custom operator, make sure you registered it with the right domain and version.

Shouldn't it be solved?

i'm facing same problem too

[philass]

@PallottaEnrico, @LebronRemonJames,

scatter_max isn't an aten torch op. Its defined in its own repository. As a result the exporters in the torch repo don't know that the op exists.

What you can do is augment the export path with a custom exporter for the scatter_max op. Note I haven't tested that the way it exports to ScatterElements is functionally correct. What I mean by this is that this will export to the ONNX ScatterElements, but scatter max is likely semantically different and you will have to adjust your export accordingly.

# POC for export custom torch ops defined outside pytorch
# Example of custom op -> ScatterMax : https://github.com/rusty1s/pytorch_scatter/blob/master/csrc/scatter.cpp#L199
#
# code below follows from instructions here https://pytorch.org/docs/stable/onnx.html#c-operators


import torch
import torch.nn as nn
from torch.onnx import symbolic_helper
from torch_scatter import scatter_max

class ScatterMax(nn.Module):

    def forward(self, src: torch.Tensor, index: torch.Tensor):
        # src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
        # index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
        out, argmax = scatter_max(src, index, dim=-1, out=src)
        return out, argmax


@symbolic_helper.parse_args("v", "v", "i", "v", "i", "i")
def symbolic_scatter_max(g, src, index, dim=-1, out=None, dim_size=None, fill_value=None):
    return (index ,g.op("ScatterElements", out, index, src, axis_i=dim, reduction_s="max"))

torch.onnx.register_custom_op_symbolic("torch_scatter::scatter_max", symbolic_scatter_max, 18)

model = ScatterMax()

src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
input_names = ["input_src", "input_idx"]
output_names = ["result_arr1", "result_arr2"]

torch.onnx.export(model, (src, index), "program.onnx", input_names=input_names, output_names=output_names, opset_version=18)

You can find the code here for reference https://github.com/philass/onnx-experiments/blob/master/export-custom.py.

[LebronRemonJames]

@PallottaEnrico, @LebronRemonJames,

scatter_max isn't an aten torch op. Its defined in its own repository. As a result the exporters in the torch repo don't know that the op exists.

What you can do is augment the export path with a custom exporter for the scatter_max op. Note I haven't tested that the way it exports to ScatterElements is functionally correct. What I mean by this is that this will export to the ONNX ScatterElements, but scatter max is likely semantically different and you will have to adjust your export accordingly.

# POC for export custom torch ops defined outside pytorch
# Example of custom op -> ScatterMax : https://github.com/rusty1s/pytorch_scatter/blob/master/csrc/scatter.cpp#L199
#
# code below follows from instructions here https://pytorch.org/docs/stable/onnx.html#c-operators


import torch
import torch.nn as nn
from torch.onnx import symbolic_helper
from torch_scatter import scatter_max

class ScatterMax(nn.Module):

    def forward(self, src: torch.Tensor, index: torch.Tensor):
        # src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
        # index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
        out, argmax = scatter_max(src, index, dim=-1, out=src)
        return out, argmax


@symbolic_helper.parse_args("v", "v", "i", "v", "i", "i")
def symbolic_scatter_max(g, src, index, dim=-1, out=None, dim_size=None, fill_value=None):
    return (index ,g.op("ScatterElements", out, index, src, axis_i=dim, reduction_s="max"))

torch.onnx.register_custom_op_symbolic("torch_scatter::scatter_max", symbolic_scatter_max, 18)

model = ScatterMax()

src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
input_names = ["input_src", "input_idx"]
output_names = ["result_arr1", "result_arr2"]

torch.onnx.export(model, (src, index), "program.onnx", input_names=input_names, output_names=output_names, opset_version=18)

You can find the code here for reference https://github.com/philass/onnx-experiments/blob/master/export-custom.py.

thanks

[paladin1410]

Hi @philass, thank you very much for your solution. I try to use your version for converting my model and it has this problem:
Traceback (most recent call last):
File "/app/LION/tools/huy_test2.py", line 91, in
torch.onnx.export(
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py", line 516, in export
_export(
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py", line 1612, in _export
graph, params_dict, torch_out = _model_to_graph(
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py", line 1138, in _model_to_graph
graph = _optimize_graph(
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py", line 677, in _optimize_graph
graph = _C._jit_pass_onnx(graph, operator_export_type)
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/utils.py", line 1956, in _run_symbolic_function
return symbolic_fn(graph_context, *inputs, **attrs)
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/symbolic_helper.py", line 395, in wrapper
return fn(g, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/symbolic_opset11.py", line 554, in cat
return opset9.cat(g, tensor_list, dim)
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/symbolic_helper.py", line 306, in wrapper
return fn(g, *args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/torch/onnx/symbolic_opset9.py", line 546, in cat
assert all(
AssertionError

Here is the code:

import torch
import torch.nn as nn
import torch_scatter
import torch.onnx
import numpy as np
from torch.onnx import symbolic_helper

class PFNLayerV2(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 use_norm=True,
                 last_layer=False):
        super().__init__()

        self.last_vfe = last_layer
        self.use_norm = use_norm
        if not self.last_vfe:
            out_channels = out_channels // 2

        if self.use_norm:
            self.linear = nn.Linear(in_channels, out_channels, bias=False)
            self.norm = nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01)
        else:
            self.linear = nn.Linear(in_channels, out_channels, bias=True)

        self.relu = nn.ReLU()

    def forward(self, inputs, unq_inv):

        x = self.linear(inputs)
        x = self.norm(x) if self.use_norm else x
        x = self.relu(x)
        x_max = torch_scatter.scatter_max(x, unq_inv, dim=0)[0]

        if self.last_vfe:
            return x_max
        else:
            x_concatenated = torch.cat([x, x_max[unq_inv, :]], dim=1)
            return x_concatenated

    
num_points = 100
in_features = 16
out_features = 32
inputs = torch.randn(num_points, in_features)
unq_inv = torch.randint(0, 10, (num_points,), dtype=torch.long)  # Assume 10 unique voxels

model = PFNLayerV2(in_features, out_features)



@symbolic_helper.parse_args("v", "v", "i", "v", "i", "i")
def symbolic_scatter_max(g, src, index, dim=-1, out=None, dim_size=None, fill_value=None):
    return (index ,g.op("ScatterElements", out, index, src, axis_i=dim, reduction_s="max"))


torch.onnx.register_custom_op_symbolic("torch_scatter::scatter_max", symbolic_scatter_max, 17)


torch.onnx.export(
    model,
    (inputs, unq_inv),
    "simple_pfn_layer.onnx",
    opset_version=17,  # Use at least opset version 16
    input_names=['inputs', 'unq_inv'],
    output_names=['output'],
)

What could be the problem? I am using pytorch 2.3.1 + cuda11.8, onnxruntime-gpu 1.18.1 and onnx 1.17.0 .Thank you