diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index f46041831525..5fdac670aa28 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1936,13 +1936,22 @@ 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) + 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 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]