This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e85fbb1fa9 [Relax][ONNX] Support scalar QDQ inputs (#20126)
e85fbb1fa9 is described below
commit e85fbb1fa93d7e417fa20d5c48925b6665daeeee
Author: Hongyi Wu <[email protected]>
AuthorDate: Wed Aug 19 07:20:39 2026 +0800
[Relax][ONNX] Support scalar QDQ inputs (#20126)
## Summary
This PR extends Relax `quantize` and `dequantize` type inference to
support
rank-zero inputs when their scale and zero point describe per-tensor
quantization.
The gap was exposed by the INT8 dynamic-shape BiRefNeXt anime
segmentation
model. Its opset-18 graph contains scalar-input `DequantizeLinear` nodes
with
scalar scale and zero-point initializers. This is valid ONNX per-tensor
dequantization, but the current Relax operator rejects the canonicalized
axis
because a rank-zero tensor has no axis.
## Goal
Import and execute valid scalar ONNX `QuantizeLinear` and
`DequantizeLinear`
graphs without weakening validation for per-axis quantization.
## What changed
- Treat the axis as irrelevant for rank-zero Relax quantize/dequantize
inputs.
- Require rank-zero inputs to use scalar or singleton scale and
zero-point
tensors.
- Keep the existing axis and parameter-size checks unchanged for inputs
with
rank one or greater.
- Add symmetric core type-inference coverage for scalar quantize and
dequantize.
- Add ONNX Runtime-backed numerical and dtype coverage for scalar
opset-18
`QuantizeLinear` and `DequantizeLinear` graphs.
## Design
ONNX uses scalar scale and zero-point tensors for per-tensor
quantization. In
that mode the `axis` attribute is ignored. The Relax ONNX frontend
already
canonicalizes the default axis to `0` for inputs with rank at most one,
but the
Relax operator type relation previously rejected every axis for a
rank-zero
input:
```text
rank(input) = 0
valid axis range = [0, rank(input) - 1] = [0, -1]
```
For rank-zero input, this change therefore skips the axis-range check
and
instead validates the condition that makes the axis irrelevant: both
scale and
zero point must be scalar or singleton tensors. The legalization and
output
shape rules already handle rank-zero tensors, so no converter or
lowering
special case is needed.
## Updated behavior
| Input / quantization parameters | Behavior |
| --- | --- |
| Rank zero + scalar/singleton scale and zero point | Accept as
per-tensor QDQ; preserve scalar output shape |
| Rank zero + non-singleton scale or zero point | Reject with an
explicit error |
| Rank one or greater | Preserve the existing axis and parameter-size
validation |
## Safety checks
- Scalar support is implemented symmetrically for quantize and
dequantize.
- A parametrized core test verifies that non-singleton per-axis
parameters are
still rejected for rank-zero input.
- The ONNX tests compare TVM execution and output dtypes with ONNX
Runtime.
- The full Relax ONNX frontend test file passes in the validation
environment,
apart from five pre-existing Float8 saturate baseline cases that were
excluded explicitly.
## Out of scope / non-goals
- Claiming full BiRefNeXt import or execution; this PR removes only its
first
observed scalar-QDQ blocker, and the model still has an independent
dynamic
shape `Concat` blocker.
- Adding the external 117 MB BiRefNeXt model to the TVM test suite.
- Changing per-axis QDQ behavior for non-scalar inputs.
- Addressing the existing Float8 saturate test baseline.
## Results
Before this change, minimal scalar ONNX graphs fail during import with:
```text
ValueError: relax.quantize: axis param is out of range (0)
ValueError: relax.dequantize: axis param is out of range (0)
```
With this change, opset-18 scalar graphs import, legalize, compile, and
execute
against ONNX Runtime:
```text
QuantizeLinear: x = 1.25, scale = 0.25, zero_point = 2 -> uint8 scalar 7
DequantizeLinear: x = 7, scale = 0.25, zero_point = 2 -> float32 scalar
1.25
```
## Tests
- `pre-commit run --files src/relax/op/tensor/qdq.cc
tests/python/relax/test_frontend_onnx.py
tests/python/relax/test_op_qdq.py`
- Fresh CPU-only H20 Release build: passed
- Relax QDQ operator tests: `9 passed`
- Focused scalar ONNX Runtime-backed tests: `2 passed, 500 deselected`
- Relax ONNX frontend H20 run: `484 passed, 9 skipped, 5 deselected, 4
xfailed`
## References
- [ONNX QuantizeLinear
specification](https://onnx.ai/onnx/operators/onnx__QuantizeLinear.html)
- [ONNX DequantizeLinear
specification](https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html)
- [BiRefNeXt anime segmentation ONNX
model](https://huggingface.co/nkta/birefnext-aniseg-ONNX)
---
src/relax/op/tensor/qdq.cc | 34 ++++++++++++++++++++++++--------
tests/python/relax/test_frontend_onnx.py | 24 ++++++++++++++++++++++
tests/python/relax/test_op_qdq.py | 32 ++++++++++++++++++++++++++++++
3 files changed, 82 insertions(+), 8 deletions(-)
diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc
index 07b9480c67..97e3f97c61 100644
--- a/src/relax/op/tensor/qdq.cc
+++ b/src/relax/op/tensor/qdq.cc
@@ -108,7 +108,7 @@ Type InferTypeQuantize(const Call& call, const
BlockBuilder& ctx) {
// Check that "axis" attribute is not out of range:
int axis = (attrs->axis < 0) ? (input_ty->ndim + attrs->axis) : attrs->axis;
- if (axis < 0 || axis > input_ty->ndim - 1) {
+ if (input_ty->ndim != 0 && (axis < 0 || axis > input_ty->ndim - 1)) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "relax.quantize: axis param is out of range (" << attrs->axis <<
")";
}
@@ -137,9 +137,18 @@ Type InferTypeQuantize(const Call& call, const
BlockBuilder& ctx) {
return false;
};
- // Check size matching of scale/zp params with input shape at dim =
attrs->axis.
- if (!is_scalar_or_singleton_vector(scale_ty)) check_param_size(scale_ty,
input_ty, "scale");
- if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty, input_ty,
"zero_point");
+ if (input_ty->ndim == 0) {
+ // A scalar input has no channel axis and only supports per-tensor
quantization.
+ if (!is_scalar_or_singleton_vector(scale_ty) ||
!is_scalar_or_singleton_vector(zp_ty)) {
+ TVM_FFI_VISIT_THROW(ValueError, call)
+ << "relax.quantize: scale and zero_point must be scalar or singleton
tensors for "
+ "rank-0 input";
+ }
+ } else {
+ // Check size matching of scale/zp params with input shape at dim =
attrs->axis.
+ if (!is_scalar_or_singleton_vector(scale_ty)) check_param_size(scale_ty,
input_ty, "scale");
+ if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty,
input_ty, "zero_point");
+ }
auto output_ty = ffi::make_object<TensorTypeNode>(*input_ty.get());
output_ty->dtype = PrimType(attrs->out_dtype);
@@ -221,7 +230,7 @@ Type InferTypeDequantize(const Call& call, const
BlockBuilder& ctx) {
// Check that "axis" attribute is not out of range:
int axis = (attrs->axis < 0) ? (input_ty->ndim + attrs->axis) : attrs->axis;
- if (axis < 0 || axis > input_ty->ndim - 1) {
+ if (input_ty->ndim != 0 && (axis < 0 || axis > input_ty->ndim - 1)) {
TVM_FFI_VISIT_THROW(ValueError, call)
<< "relax.dequantize: axis param is out of range (" << attrs->axis <<
")";
}
@@ -250,9 +259,18 @@ Type InferTypeDequantize(const Call& call, const
BlockBuilder& ctx) {
return false;
};
- // Check size matching of scale/zp params with input shape at dim =
attrs->axis.
- if (!is_scalar_or_singleton_vector(scale_ty)) check_param_size(scale_ty,
input_ty, "scale");
- if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty, input_ty,
"zero_point");
+ if (input_ty->ndim == 0) {
+ // A scalar input has no channel axis and only supports per-tensor
quantization.
+ if (!is_scalar_or_singleton_vector(scale_ty) ||
!is_scalar_or_singleton_vector(zp_ty)) {
+ TVM_FFI_VISIT_THROW(ValueError, call)
+ << "relax.dequantize: scale and zero_point must be scalar or
singleton tensors for "
+ "rank-0 input";
+ }
+ } else {
+ // Check size matching of scale/zp params with input shape at dim =
attrs->axis.
+ if (!is_scalar_or_singleton_vector(scale_ty)) check_param_size(scale_ty,
input_ty, "scale");
+ if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty,
input_ty, "zero_point");
+ }
auto output_ty = ffi::make_object<TensorTypeNode>(*input_ty.get());
output_ty->dtype = PrimType(attrs->out_dtype);
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 32dd4b0bef..673c25bcdd 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -12476,6 +12476,30 @@ def test_dequantizelinear_singleton_qparams_opset10():
check_correctness(model, inputs={"x": x}, opset=10, check_dtypes=True)
[email protected](
+ ("op_type", "input_dtype", "output_dtype", "input_value"),
+ [
+ ("QuantizeLinear", TensorProto.FLOAT, TensorProto.UINT8,
np.array(1.25, "float32")),
+ ("DequantizeLinear", TensorProto.UINT8, TensorProto.FLOAT, np.array(7,
"uint8")),
+ ],
+)
+def test_qdqlinear_scalar_input(op_type, input_dtype, output_dtype,
input_value):
+ node = helper.make_node(op_type, ["x", "scale", "zero_point"], ["y"])
+ graph = helper.make_graph(
+ [node],
+ f"{op_type.lower()}_scalar_input",
+ [helper.make_tensor_value_info("x", input_dtype, [])],
+ [helper.make_tensor_value_info("y", output_dtype, [])],
+ initializer=[
+ helper.make_tensor("scale", TensorProto.FLOAT, [], [0.25]),
+ helper.make_tensor("zero_point", TensorProto.UINT8, [], [2]),
+ ],
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
18)])
+
+ check_correctness(model, inputs={"x": input_value}, opset=18,
check_dtypes=True)
+
+
def test_quantizelinear_optional_zero_point_opset13():
"""ONNX allows missing zero_point input; importer should default it to 0
(uint8)."""
node = helper.make_node("QuantizeLinear", ["x", "scale"], ["y"])
diff --git a/tests/python/relax/test_op_qdq.py
b/tests/python/relax/test_op_qdq.py
index 960dacc673..6ad35b6e85 100644
--- a/tests/python/relax/test_op_qdq.py
+++ b/tests/python/relax/test_op_qdq.py
@@ -14,6 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+import pytest
+
import tvm
import tvm.testing
from tvm import relax, tirx
@@ -49,6 +51,36 @@ def test_qdq_op_infer_ty():
)
+def test_qdq_op_infer_ty_scalar():
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((), "float32"))
+ dx = relax.Var("dx", R.Tensor((), "uint8"))
+ s = relax.Var("s", R.Tensor((), "float32"))
+ zp = relax.Var("zp", R.Tensor((), "uint8"))
+ _check_inference(bb, relax.op.quantize(x, s, zp, 0, "uint8"),
relax.TensorType((), "uint8"))
+ _check_inference(
+ bb,
+ relax.op.dequantize(dx, s, zp, 0, "float32"),
+ relax.TensorType((), "float32"),
+ )
+
+
[email protected]("op_name", ["quantize", "dequantize"])
+def test_qdq_op_rejects_per_axis_params_for_scalar_input(op_name):
+ bb = relax.BlockBuilder()
+ s = relax.Var("s", R.Tensor((2,), "float32"))
+ zp = relax.Var("zp", R.Tensor((2,), "uint8"))
+ if op_name == "quantize":
+ x = relax.Var("x", R.Tensor((), "float32"))
+ call = relax.op.quantize(x, s, zp, 0, "uint8")
+ else:
+ x = relax.Var("x", R.Tensor((), "uint8"))
+ call = relax.op.dequantize(x, s, zp, 0, "float32")
+
+ with pytest.raises(ValueError, match="scalar or singleton tensors"):
+ bb.normalize(call)
+
+
def test_qdq_op_infer_ty_unknown_dtype():
bb = relax.BlockBuilder()
x = relax.Var("x", R.Tensor((2, 3), dtype=None))
