[torchlib] Implement quantize_per_channel and dequantize_per_channel

[Copilot]

This PR implements the missing quantize_per_channel and dequantize_per_channel operations in the torchlib quantized_decomposed module.

Changes

Added two new functions to onnxscript/function_libs/torch_lib/ops/quantized_decomposed.py:

quantized_decomposed_quantize_per_channel

• Implements per-channel quantization using ONNX QuantizeLinear with per-axis support
• Takes tensor inputs for scales and zero_points (one value per channel)
• Supports axis parameter to specify the quantization dimension
• Uses ONNX opset23 for per-axis quantization capabilities

quantized_decomposed_dequantize_per_channel

• Implements per-channel dequantization using ONNX DequantizeLinear with per-axis support
• Takes tensor inputs for scales and optional zero_points
• zero_points parameter is Optional[TensorType] matching PyTorch reference
• Supports both default output type and explicit output_dtype parameter

Implementation Details

Both functions:

• Follow the existing code patterns using @torch_op decorator with trace_only=True
• Have function signatures that match the PyTorch reference implementations from torch.ao.quantization.fx._decomposed
• Use ONNX opset23 to leverage the axis and output_dtype parameters for per-axis quantization
• Are properly registered and accessible as TracedOnnxFunctions

The implementation leverages ONNX's native per-axis quantization support rather than implementing the tensor manipulation logic from the PyTorch reference, making it more efficient and aligned with ONNX best practices.

Testing

Validated that:

• Functions import correctly and have proper signatures
• Function schemas are correctly generated with appropriate parameter types
• No existing functionality is broken (per-tensor functions still work)
• Functions can be used to build ONNX models

Fixes #2389.

---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

[justinchuby]

Are these unused? Why is that?

[Copilot]

These parameters are unused because the ONNX QuantizeLinear/DequantizeLinear operators don't require explicit quant_min/quant_max parameters - they determine the quantization range from the data type and quantization parameters. The parameters are kept in the function signature for API compatibility with PyTorch's reference implementation, following the same pattern as the existing per-tensor functions above.

[codecov]

❌ 3 Tests Failed:

Tests completed	Failed	Passed	Skipped
15728	3	15725	1701

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0756_test_or_bcast3v2d

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.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_or_bcast3v2d'

The above exception was the direct cause of the following exception:
.nox\test\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_or_bcast3v2d' (e=No module named 'tests.onnx_backend_test_code.test_or_bcast3v2d') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_or_bcast3v2d.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_or_bcast3v2d.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
E   from onnxscript.onnx_opset import opset7
E   
E   @script()
E   def bck_test_or_bcast3v2d(x: BOOL[3,4,5], y: BOOL[4,5]) -> (BOOL[3,4,5]):
E       r_or = opset7.Or(x, y)
E       return r_or

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1100_test_shape_end_1

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.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_shape_end_1'

The above exception was the direct cause of the following exception:
.nox\test\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_shape_end_1' (e=No module named 'tests.onnx_backend_test_code.test_shape_end_1') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_1.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_1.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 FLOAT, INT64
E   from onnxscript.onnx_opset import opset21
E   
E   @script()
E   def bck_test_shape_end_1(x: FLOAT[3,4,5]) -> (INT64[1]):
E       y = opset21.Shape(x, end=1)
E       return y

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0805_test_reduce_l1_do_not_keepdims_example

Stack Traces | 0.004s 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_reduce_l1_do_not_keepdims_example'

The above exception was the direct cause of the following exception:
.nox\test\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_reduce_l1_do_not_keepdims_example' (e=No module named 'tests.onnx_backend_test_code.test_reduce_l1_do_not_keepdims_example') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_do_not_keepdims_example.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_reduce_l1_do_not_keepdims_example.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 FLOAT, INT64
E   from onnxscript.onnx_opset import opset18
E   
E   @script()
E   def bck_test_reduce_l1_do_not_keepdims_example(data: FLOAT[3,2,2], axes: INT64[1]) -> (FLOAT[3,2]):
E       reduced = opset18.ReduceL1(data, axes, keepdims=0)
E       return reduced

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]

Please format the code using lintrunner. You can do

pip install --upgrade lintrunner lintrunner-adapters
lintrunner init

then

lintrunner f --all-files

[Copilot]

Please format the code using lintrunner. You can do

...

Done! Code has been formatted using lintrunner. The formatting fixed import sorting, whitespace issues, and line wrapping. Commit: 82c8f9e