From 6c4678140692869f5a160b5e457e17a797b407f0 Mon Sep 17 00:00:00 2001 From: OmarAzizi Date: Tue, 7 Jul 2026 18:49:24 +0300 Subject: [PATCH 1/2] [Fix] Preserve ONNX Squeeze axes attribute for opset < 13 --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 10 ++++++ tests/python/relax/test_frontend_onnx.py | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index f46041831525..990a0583168b 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1936,13 +1936,23 @@ def _impl_v14(cls, bb, inputs, attr, params): 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) + axis = tuple(int(x) for x in axes) if axes is not None else None + return cls._squeeze(bb, inputs[0], axis) + @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 9058cdf70ae7..915a522be88d 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -3729,6 +3729,42 @@ def main(x: R.Tensor((1, 32, 1, 32), dtype="float32")) -> R.Tensor( 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] From f129929df47e715865c5c378c1cc1efedc6655f2 Mon Sep 17 00:00:00 2001 From: OmarAzizi Date: Sat, 11 Jul 2026 23:16:25 +0300 Subject: [PATCH 2/2] Remove redundant axes tuple conversion in Squeeze._impl_v1 --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 990a0583168b..5fdac670aa28 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1940,8 +1940,7 @@ class Squeeze(OnnxOpConverter): 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) - axis = tuple(int(x) for x in axes) if axes is not None else None - return cls._squeeze(bb, inputs[0], axis) + return cls._squeeze(bb, inputs[0], axes) @classmethod def _impl_v13(cls, bb, inputs, attr, params):