Skip to content

[torchlib] Implement aten__upsample_bicubic2d_aa and aten__upsample_bilinear2d_aa functions #2383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 17, 2025

Conversation

Copilot
Copy link
Contributor

@Copilot Copilot AI commented Jun 14, 2025

This PR implements the missing anti-aliasing (AA) variants of upsample functions that were requested in issue #1159:

  • aten__upsample_bicubic2d_aa - bicubic 2D upsampling with anti-aliasing
  • aten__upsample_bilinear2d_aa - bilinear 2D upsampling with anti-aliasing

Changes Made

Core Implementation

  • Modified helper functions to support anti-aliasing:

    • Added antialias parameter (default=0) to _aten_upsample_output_size()
    • Added antialias parameter (default=0) to _aten_upsample_scales()
    • Maintains backward compatibility with existing code
  • Implemented AA functions with same signatures as regular variants:

    def aten__upsample_bicubic2d_aa(self, output_size, align_corners, scales_h=None, scales_w=None)
    def aten__upsample_bilinear2d_aa(self, output_size, align_corners, scales_h=None, scales_w=None) 

    Both functions pass antialias=1 to enable ONNX Resize anti-aliasing.

Test Configuration

  • Added OpInfo entries in extra_opinfo.py for both AA functions
  • Added TorchLibOpInfo entries in ops_test_data.py with compare_shape_only_for_output=(0,) since ONNX and PyTorch use different anti-aliasing algorithms

Technical Details

The AA variants use the same underlying logic as regular upsample functions but enable anti-aliasing in the ONNX Resize operation. As noted in the original issue discussion, ONNX and PyTorch implement different anti-aliasing methods, so tests compare shapes rather than exact values.

Example usage:

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

# Create test input
input_tensor = np.array([[[[2,1,1,1], [1,1,1,1], [1,1,1,1], [1,1,1,1]]]]).astype(np.float32)
output_size = np.array([1,1]).astype(np.int64)

# Use AA upsampling
result = aten__upsample_bicubic2d_aa(input_tensor, output_size, align_corners=True)
print(result)  # Output: [[[[1.390625]]]]

Testing Results

  • ✅ All new AA function tests pass (2 passed, 1 skipped as expected for trace-only functions)
  • ✅ All existing upsample function tests continue to pass - no regressions
  • ✅ Functions produce expected different output when AA is enabled vs disabled
  • ✅ Helper functions work correctly with both antialias=0 and antialias=1

This implementation matches the approach from the previous PR #1259 and completes the upsample function suite requested in the issue.

Fixes #1159. Fixes pytorch/pytorch#128818


💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

@Copilot Copilot AI changed the title [WIP] [torchlib] Implement various upsample functions [torchlib] Implement aten__upsample_bicubic2d_aa and aten__upsample_bilinear2d_aa functions Jun 14, 2025
@Copilot Copilot AI requested a review from justinchuby June 14, 2025 03:59
Copilot finished work on behalf of justinchuby June 14, 2025 03:59
Copy link

codecov bot commented Jun 14, 2025

❌ 4 Tests Failed:

Tests completed Failed Passed Skipped
16447 4 16443 2362
View the top 3 failed test(s) by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0740_test_max_float64
Stack Traces | 0.003s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_max_float64'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_max_float64' (e=No module named 'tests.onnx_backend_test_code.test_max_float64') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_max_float64.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_max_float64.py', current folder: C:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import DOUBLE
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_max_float64(data_0: DOUBLE[3], data_1: DOUBLE[3]) -> (DOUBLE[3]):
E       result = opset13.Max(data_0, data_1)
E       return result
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0360_test_compress_0
Stack Traces | 0.007s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_compress_0'

The above exception was the direct cause of the following exception:
.nox\test_onnx_weekly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_compress_0' (e=No module named 'tests.onnx_backend_test_code.test_compress_0') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_compress_0.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_compress_0.py', current folder: C:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import BOOL, FLOAT
E   from onnxscript.onnx_opset import opset11
E   
E   @script()
E   def bck_test_compress_0(input: FLOAT[3,2], condition: BOOL[3]) -> (FLOAT[2,2]):
E       output = opset11.Compress(input, condition, axis=0)
E       return output
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0904_test_ai_onnx_ml_tree_ensemble_set_membership
Stack Traces | 0.009s run time
onnxscript/converter.py:460: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../hostedtoolcache/Python/3.11.12.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:186: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
     ^^^^^^^^
onnxscript/main.py:94: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/main.py:38: in script_check
    return convert.translate_function_def(f)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:1452: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:1439: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:961: in _translate_stmt
    return self._translate_assign_stmt(node)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:1048: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:992: in assign
    t = self._translate_expr(rhs, lhs).name
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:546: in _translate_expr
    r = self._translate_call_expr(node)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:825: in _translate_call_expr
    attrs = [
onnxscript/converter.py:826: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:510: in _translate_attr
    val = self._eval_constant_expr(expr)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/converter.py:462: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

To view more test analytics, go to the Test Analytics Dashboard
📋 Got 3 mins? Take this short survey to help us improve Test Analytics.

@justinchuby justinchuby marked this pull request as ready for review June 14, 2025 15:11
@justinchuby justinchuby enabled auto-merge (squash) June 14, 2025 15:11
@justinchuby justinchuby requested a review from titaiwangms June 14, 2025 15:11
Copilot finished work on behalf of titaiwangms June 16, 2025 20:06
Copy link
Contributor

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested with pytorch/pytorch#128818, and it works!

@titaiwangms titaiwangms disabled auto-merge June 17, 2025 19:35
@titaiwangms titaiwangms merged commit 0582b6b into main Jun 17, 2025
28 of 32 checks passed
@justinchuby justinchuby deleted the copilot/fix-1159 branch June 17, 2025 19:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
Development

Successfully merging this pull request may close these issues.

[torchlib] Implement various upsample functions ONNX Dynamo Export - Unsupported FX nodes: {'call_function': ['aten._upsample_bilinear2d_aa.default']}.
3 participants