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 396dd34946 [Fix][Relax][ONNX] Preserve ONNX Squeeze axes attribute for
opset < 13 (#19966)
396dd34946 is described below
commit 396dd349460ffbfd1e46d9308ca7d67997d0ad23
Author: Kryptonite <[email protected]>
AuthorDate: Sat Jul 18 07:14:55 2026 +0300
[Fix][Relax][ONNX] Preserve ONNX Squeeze axes attribute for opset < 13
(#19966)
## Summary
Before opset 13, ONNX `Squeeze` specifies `axes` as a node attribute
rather than a tensor input. The Relax ONNX importer only implemented
`_impl_v13`, which reads axes from the second input, so for opset < 13
models, the attribute was silently ignored (`axis` defaulted to `None`)
and the importer squeezed every size-1 dimension instead of only the
requested one. This produced tensors with the wrong rank, breaking
downstream ops like `Transpose` whose `perm` no longer matched the
input's actual rank.
Added `_impl_v1` to read `axes` from the node attribute for opset < 13,
and factored the existing squeeze logic into a shared `_squeeze` helper
used by both `_impl_v1` and `_impl_v13`.
## Test plan
- Added `test_squeeze_axes_attribute` to
`tests/python/relax/test_frontend_onnx.py`, covering an opset-11
`Squeeze` node with `axes` as an attribute.
- Ran `pytest tests/python/relax/test_frontend_onnx.py -k squeeze`. All
21 tests pass.
- Verified against the real-world model that triggers this bug,
[PaddlePaddle/PP-OCRv6_tiny_rec_onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_tiny_rec_onnx)
(opset 11, uses attribute-based `Squeeze`): import fails on `main` with
`Transpose: number of axes in perm attribute (3) must equal the number
of input tensor dimensions (-1)`, and succeeds with this fix.
## Real-world reproduction
```python
import urllib.request
import onnx
from tvm.relax.frontend.onnx import from_onnx
# PaddlePaddle/PP-OCRv6_tiny_rec_onnx (opset 11, uses attribute-based
Squeeze)
url =
"https://huggingface.co/PaddlePaddle/PP-OCRv6_tiny_rec_onnx/resolve/main/inference.onnx";
path = "pp_ocrv6_tiny_rec.onnx"
urllib.request.urlretrieve(url, path)
model = onnx.load(path)
print("opset:", [(o.domain, o.version) for o in model.opset_import])
for node in model.graph.node:
if node.op_type == "Squeeze":
axes_attr = [a for a in node.attribute if a.name == "axes"]
print(node.name, "inputs=", list(node.input), "axes_attr=",
axes_attr)
# Fails on main with:
# ValueError: Transpose: number of axes in perm attribute (3) must equal
the number of input tensor dimensions (-1)
# Succeeds with this fix.
mod = from_onnx(model)
print("Import succeeded")
```
Fixes (partially) #19965. The shape-Gather and dynamic-TopK issues
reported in that issue are separate and not addressed here.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 9 +++++++
tests/python/relax/test_frontend_onnx.py | 36 +++++++++++++++++++++++++
2 files changed, 45 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index eaad127524..9121018bb9 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2087,13 +2087,22 @@ class CumSum(OnnxOpConverter):
class Squeeze(OnnxOpConverter):
"""Converts an onnx Squeeze node into an equivalent Relax expression."""
+ @classmethod
+ def _impl_v1(cls, bb, inputs, attr, params):
+ # Prior to opset 13, axes is provided as an attribute rather than an
input.
+ axes = attr.get("axes", None)
+ return cls._squeeze(bb, inputs[0], axes)
+
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
data = inputs[0]
axis = get_constant(inputs[1], params)
if isinstance(axis, relax.Constant):
axis = tuple([int(x) for x in axis.data.numpy()])
+ return cls._squeeze(bb, data, axis)
+ @classmethod
+ def _squeeze(cls, bb, data, axis):
# If data is constant, perform computation directly.
if isinstance(data, relax.Constant):
if isinstance(axis, tuple | type(None)):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index db9bf18a81..4a5c2f778d 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3687,6 +3687,42 @@ def test_squeeze():
verify_squeeze(None, ExpectedSqueezeAll)
+def test_squeeze_axes_attribute():
+ # Prior to opset 13, ONNX Squeeze takes `axes` as an attribute rather than
an input.
+ squeeze_node = helper.make_node("Squeeze", ["x"], ["y"], axes=[0, 2])
+ shape = [1, 32, 1, 32]
+
+ graph = helper.make_graph(
+ [squeeze_node],
+ "squeeze_axes_attribute_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32,
32])],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="squeeze_axes_attribute_test",
+ opset_imports=[helper.make_opsetid("", 11)],
+ )
+ tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 32, 1, 32), dtype="float32")) -> R.Tensor(
+ (32, 32), dtype="float32"
+ ):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((32, 32), dtype="float32") = R.squeeze(x,
axis=[0, 2])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
def test_squeeze_constant():
def verify_squeeze_constant(axis, expected):
shape = [1, 2, 1, 3]
