[Fix][Relax][ONNX] Preserve ONNX Squeeze axes attribute for opset < 13

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)):

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]