Constant Folding removes 'axes' initializer form Squeeze layer, resulting in incorrect ONNX model

[bas-aarts]

converting the attached ONNX model with the following script:

from onnxscript import ir
from onnxscript.optimizer import optimize

ir_model = ir.load("input.onnx")
optimized_model = optimize(ir_model)
ir.save(optimized_model, "output.onnx")

produces an incorrect ONNX model.
axes initializer '17' in Squeeze layer 'If_11_Squeeze_13' is missing.
The constant folder is causing this.

The If layer in the input model is the default way that a squeeze layer is exported from Torch, making this a common pattern.

input.onnx.zip

versions of relevant packages:

ml_dtypes         0.5.4
onnx              1.20.1
onnx-ir           0.1.16
onnxscript        0.6.2

[justinchuby]

Can you elaborate what is incorrect and which program raises the error? Squeeze's axes input is optional according to https://onnx.ai/onnx/operators/onnx__Squeeze.html

[bas-aarts]

like I said, the initializer is missing:

Python 3.10.12 (main, May 27 2025, 17:12:29) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import onnx
>>> onnx.checker.check_model("output.onnx", full_check=True)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/local/mnt/workspace/git/tetracode/tetra-dev/lib/python3.10/site-packages/onnx/checker.py", line 164, in check_model
    C.check_model_path(
onnx.onnx_cpp2py_export.checker.ValidationError: Nodes in a graph must be topologically sorted, however input '17' of node:
name: If_11_Squeeze_13 OpType: Squeeze
 is not output of any previous nodes.
>>>

[justinchuby]

I see the issue now. The constant folder for the If node omits initializers from the subgraphs.

@bas-aarts for now you can call https://onnx.ai/ir-py/api/ir_passes_common.html#onnx_ir.passes.common.LiftSubgraphInitializersToMainGraphPass on the model to move the initializers (constant values attached to the graph) before optimizing.

Also note that the model is using opset 13. The optimizer only supports opset version >= 18. You may use the onnx version converter to convert the model to a newer version first. (that's unrelated to bug reported in this issue).

cc @titaiwangms

[bas-aarts]

Awesome, that works. Thanks!