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 c9e9b85c14 [Fix][Relax] Raise error on non-unit dim ONNX Squeeze axis
(#20188)
c9e9b85c14 is described below
commit c9e9b85c14a19cfe9ce1880c6ecafefa4729cb3e
Author: Kryptonite <[email protected]>
AuthorDate: Fri Aug 28 01:01:38 2026 +0300
[Fix][Relax] Raise error on non-unit dim ONNX Squeeze axis (#20188)
TVM's ONNX Squeeze converter accepted axes with non-unit dimensions
instead of raising an error silently passing tensors through unchanged.
This adds a check in the ONNX frontend that raises ValueError when a
squeeze axis's statically known size isn't 1, matching the ONNX spec.
Fixes #20185
Edit: The code merged from #20147 was causing a linter error so Claude
assisted me in reformatting the code
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 23 ++++++++++++
tests/python/relax/test_frontend_onnx.py | 48 +++++++++++++++++++++++++
2 files changed, 71 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index d078bdf9f6..fdcf26d170 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2243,6 +2243,28 @@ class Squeeze(OnnxOpConverter):
axis = tuple([int(x) for x in axis.data.numpy()])
return cls._squeeze(bb, data, axis)
+ @classmethod
+ def _check_squeeze_axes_are_unit_dims(cls, data, axes):
+ """Raise if any axis to be squeezed has a statically known extent
other than 1."""
+ rank = _get_known_tensor_rank(data)
+ if rank is None:
+ return
+ ty = data.ty
+ if not (isinstance(ty, relax.TensorType) and isinstance(ty.shape,
relax.ShapeExpr)):
+ return
+ for axis in _normalize_constant_axes(list(axes), rank, "Squeeze"):
+ extent = ty.shape.values[axis]
+ if not isinstance(extent, tirx.IntImm):
+ raise ValueError(
+ f"Squeeze axis {axis} has a symbolic extent that cannot be
proven to be "
+ "1 at import time; only statically known unit-size axes
can be squeezed."
+ )
+ if int(extent.value) != 1:
+ raise ValueError(
+ f"Squeeze axis {axis} has size {int(extent.value)}, but
only "
+ "axes of size 1 can be squeezed."
+ )
+
@classmethod
def _squeeze(cls, bb, data, axis):
# If data is constant, perform computation directly.
@@ -2271,6 +2293,7 @@ class Squeeze(OnnxOpConverter):
return relax.op.squeeze(data)
if isinstance(axis, tuple):
+ cls._check_squeeze_axes_are_unit_dims(data, axis)
return relax.op.squeeze(data, list(axis))
data_ndim = _get_known_tensor_rank(data)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 4e31015b10..6a417d56f1 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -4547,6 +4547,54 @@ def test_squeeze_axes_attribute():
tvm.ir.assert_structural_equal(tvm_model, Expected)
+def test_squeeze_non_unit_axis_raises():
+ # Per the ONNX spec, squeezing an axis whose length is not 1 must raise an
error
+ squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
+ shape = [2, 3]
+
+ graph = helper.make_graph(
+ [squeeze_node],
+ "squeeze_non_unit_axis_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
+ ],
+ initializer=[helper.make_tensor("axes", TensorProto.INT64, [1], [0])],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3])],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="squeeze_non_unit_axis_test",
+ opset_imports=[helper.make_opsetid("", 13)],
+ )
+ with pytest.raises(ValueError, match="has size 2"):
+ from_onnx(model, opset=13, keep_params_in_input=True)
+
+
+def test_squeeze_symbolic_axis_raises():
+ # Squeezing an axis whose extent is symbolic can't be proven to be 1 at
import time
+ squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
+ shape = ["N", 3]
+
+ graph = helper.make_graph(
+ [squeeze_node],
+ "squeeze_symbolic_axis_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
+ ],
+ initializer=[helper.make_tensor("axes", TensorProto.INT64, [1], [0])],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3])],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="squeeze_symbolic_axis_test",
+ opset_imports=[helper.make_opsetid("", 13)],
+ )
+ with pytest.raises(ValueError, match="symbolic extent"):
+ from_onnx(model, opset=13, keep_params_in_input=True)
+
+
def test_squeeze_constant():
def verify_squeeze_constant(axis, expected):
shape = [1, 2, 1, 3]
