Exporting a model with MultiheadAttention and dynamic shapes fails

[f-fuchs]

test_reshape_minimal.py

onnx_export_option_B_pt_export.md
onnx_export_option_A_success.md

🐛 Describe the bug

Hey,

I am trying to export a PyTorch UNet model with MultiheadAttention in the bottleneck and dynamic shapes.
Without MultiheadAttention I can use dynamo=False and inference is successful with variable input shapes but this fails with MultiheadAttention .

In this gist I tried exporting with dynamo and dynamic_axes, with dynamo and dynamic_shapes and without dynamo but every option fails either during exporting or during inference.

https://gist.github.com/f-fuchs/012f7efb2984700d2eca93e5a8acc3f3

Option A: dynamo with dynamic_axes

export succeeds but inference with but size =! 1 leads to:

onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: input for the following indices
index: 0 Got: 2 Expected: 1
Please fix either the inputs/outputs or the model.

Option B: dynamo with dynamic_shapes

export fails with convoluted guards, report is attached

Option C: dynamo=False with dynamic_axes

export works but error during inference due to MultiheadAttention
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/attention_bottleneck/transformer_blocks/transformer_blocks.0/attn_self/Reshape_4' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{1024,1,512}, requested shape:{256,8,64}

Versions

torch 2.7.1
onnx 1.18.0

[justinchuby]

Please try using a dim size that is >=2 (not 0 or 1) in the example input: The 0/1 specialization problem in Pt2 Export and podcast on this topic

Also please use the latest PyTorch version and let us know. Thanks!

[f-fuchs]

After updating to torch 2.8.0 the behavior of option C stayed untouched. But A and B are both failing to export now but with different errors, see the attached outputs and reports.

onnx_export_2025-10-10_09-10-26-269056_conversion_optionB.md
output_optionB.txt
onnx_export_2025-10-10_09-09-35-673291_decomp_optionA.md
output_optionA.txt

[justinchuby]

The error log states that there is a missing op sym_size which should be available: https://github.com/microsoft/onnxscript/blob/59c3d32ea0cf18fbd348d8b4e23fdb8dad6427ea/onnxscript/function_libs/torch_lib/ops/core.py#L8216

Could you ensure onnxscript is up to date, and you are using an opset version >= 18?

[f-fuchs]

Thx setting the opset version to 18 worked. Option A and B are now both working at least as long as the input data has a batch size > 1.