Add op (AvgPool) | feat (torchlib)

[titaiwangms]

1D and 3D

Fix can save the op after opset 19, but not before it. So an xfail is created when ceil_mode=True

[titaiwangms]

I am blocked by ceil_mode. It seems that the ceil_mode in ORT is not following ONNX spec:

import numpy as np
from onnxscript.function_libs.torch_lib.ops.nn import aten_avg_pool2d

x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)
kernal = [3, 3]
stride = [2, 2]

print(aten_avg_pool2d(x, kernel_size=kernal, stride=stride, ceil_mode=True))
# Tensor(array([[[[6., 5.], [8., 6.]]]], dtype=float32))

# https://github.com/onnx/onnx/blob/main/docs/Operators.md#averagepool
# The answer should be:
# np.array([[[[6, 7.5], [12, 13.5]]]]).astype(np.float32)

So the difference is that in ceil_mode, ORT tends to divide the value by kernal_size no matter the right padding exists or not.
But I don't think they would take this as repro..

[titaiwangms]

I am blocked by ceil_mode. It seems that the ceil_mode in ORT is not following ONNX spec:

import numpy as np
from onnxscript.function_libs.torch_lib.ops.nn import aten_avg_pool2d

x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)
kernal = [3, 3]
stride = [2, 2]

print(aten_avg_pool2d(x, kernel_size=kernal, stride=stride, ceil_mode=True))
# Tensor(array([[[[6., 5.], [8., 6.]]]], dtype=float32))

# https://github.com/onnx/onnx/blob/main/docs/Operators.md#averagepool
# The answer should be:
# np.array([[[[6, 7.5], [12, 13.5]]]]).astype(np.float32)

So the difference is that in ceil_mode, ORT tends to divide the value by kernal_size no matter the right padding exists or not. But I don't think they would take this as repro..

Just found this issue is caused by "count_include_pad=True will pad the right side even if it's not padded". So the following code would be able to solve this:

x = np.array(
    [
        [
            [
                [1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12],
                [13, 14, 15, 16],
            ]
        ]
    ]
).astype(np.float32)

import onnxscript
from onnxscript.onnx_opset import opset18 as op

@onnxscript.script(default_opset=op)
def avg_pool(x):
    result = op.AveragePool(x, kernel_shape=[3,3], strides=[2,2], ceil_mode=True, count_include_pad=False)
    return result

print(avg_pool(x))

Although I still think this doesn't make sense.

[titaiwangms]

However, the above doesn't solve the root cause which is that in ORT, the last slide of window tends to dividing the value with kernal size regardless the remain pixel+pads seems like a bug.

[titaiwangms]

Not sure if it's a spec issue or implementation: onnx/onnx#5276

[titaiwangms]

microsoft/onnxruntime#16203

[codecov]

Codecov Report

Merging #754 (5b5b715) into main (65d7c03) will decrease coverage by 0.04%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #754      +/-   ##
==========================================
- Coverage   76.67%   76.64%   -0.04%     
==========================================
  Files         112      112              
  Lines       13496    13508      +12     
  Branches     1365     1366       +1     
==========================================
+ Hits        10348    10353       +5     
- Misses       2809     2817       +8     
+ Partials      339      338       -1     

Impacted Files	Coverage Δ
onnxscript/function_libs/torch_lib/ops/core.py	77.10% <ø> (+0.02%)	⬆️
onnxscript/function_libs/torch_lib/ops/nn.py	71.42% <100.00%> (+1.08%)	⬆️
...ipt/tests/function_libs/torch_lib/ops_test_data.py	96.78% <100.00%> (+0.01%)	⬆️

... and 1 file with indirect coverage changes