Your ONNX model says its inputs are dynamic. Some of them aren't
When you export with dynamic_shapes, the ONNX model's input axes come out symbolic, and it is natural to assume the model now works for any size. I tested that assumption op by op: export each op with every dimension dynamic, then run it at a different size and compare with PyTorch. Most ops passed. Two did not, and one of them breaks torchvision's VGG. I reported it as pytorch/pytorch#204405.
The test
About 60 ops, each in a small module, exported with torch.export.Dim.AUTO on every axis. Each was then run on ONNX Runtime with inputs of a new size, for example exported at (2, 5) and run at (3, 9), and compared with eager PyTorch. Slicing, padding, pooling, interpolation, normalization, sorting, attention, STFT and the rest all handled the new sizes. adaptive_avg_pool2d and kron did not.
adaptive_avg_pool2d: right at one size, wrong at every other
class M(torch.nn.Module):
def forward(self, x):
return F.adaptive_avg_pool2d(x, (2, 3))
x = torch.randn(1, 2, 6, 6)
dyn = ({0: Dim.AUTO, 1: Dim.AUTO, 2: Dim.AUTO, 3: Dim.AUTO},)
onnx_program = torch.onnx.export(M(), (x,), dynamic_shapes=dyn, dynamo=True)
| Input | PyTorch | ONNX |
|---|---|---|
| (1, 2, 6, 6), the example | (1, 2, 2, 3) | (1, 2, 2, 3) |
| (2, 3, 8, 9) | (2, 3, 2, 3) | (2, 3, 2, 4) |
| (1, 1, 12, 12) | (1, 1, 2, 3) | (1, 1, 4, 6) |
There is no error and no warning. The ONNX graph is one AveragePool with kernel_shape=(3, 2) and strides=(3, 2), computed from the 6x6 example, while the input axes are still declared dynamic.
This is not an exotic op. torchvision's vgg11() ends its backbone with AdaptiveAvgPool2d((7, 7)). Exported with dynamic height and width, it works at 224x224. At 256x256 and 448x448, ONNX Runtime fails in the classifier's Gemm. Without the classifier, the backbone returns (2, 512, 8, 8) and (2, 512, 14, 14) instead of (2, 512, 7, 7), silently.
Where the size gets baked in
torch.export itself is fine. It keeps aten._adaptive_avg_pool2d as one op, and the exported program gives the right shape, (1, 1, 2, 3), for a 12x12 input. The problem is a later step in the ONNX exporter:
- torchlib, which maps ATen ops to ONNX, has no
_adaptive_avg_pool2d. It was removed because ONNX has no general adaptive pooling op, so the exporter falls back to PyTorch's own decomposition. - That decomposition has a fast path: if the input size divides the output size, call
avg_pool2dwithkernel = input // output. With symbolic sizes, this evaluates the condition and the kernel for the example input, and the decomposed graph becomesavg_pool2d(x, [3, 2], [3, 2]). - This happens after export, during the exporter's own decomposition, so the program's range constraints still say the dims are dynamic, and so does the ONNX model.
kron follows the same pattern. It is decomposed into _unsafe_view calls with the example's sizes baked in, and at a new size ONNX Runtime fails in a Reshape. There is one useful detail: the decomposed ExportedProgram does record the specialization. Running it at the new size raises Guard failed: a.size()[0] == 2. So the exporter could detect this case by comparing guards before and after its decomposition, and warn instead of writing a model that is wrong for other sizes.
One false alarm
conv1d(padding="same") also failed to export at first. I had made the weight dynamic too, so the padding depended on a symbolic kernel size, which an ONNX Conv attribute cannot hold. In a real model the weight is a parameter with a fixed shape. With only the input dynamic, Conv1d(..., padding="same") exports and matches PyTorch at new sizes for every kernel size and dilation I tried. Make your test match how the op is used in real models before you file a bug.
What to do today
- After exporting with dynamic shapes, run the ONNX model at two or three sizes other than the example and compare with PyTorch. One run at the example size proves nothing about dynamic axes.
- If your model uses
AdaptiveAvgPool2dwith an output size other than 1, either export with static spatial sizes or check the exportedAveragePoolattributes. An output size of 1 passed my test. - Follow pytorch/pytorch#204405 for the fix. I proposed a shape-dynamic torchlib implementation, built from
CumSumandGather, and a general check in the exporter. The implementation is now up for review as microsoft/onnxscript#3201. With it, VGG11 exported with dynamic height and width matches PyTorch at 224, 256 and 448.
I used Claude (Anthropic) to help build the tests, investigate the findings and draft this post. I reviewed the results and ran the checks myself.
Qiming Bao is an AI researcher in Auckland, New Zealand, working on LLM reasoning and model deployment. He received his PhD from the University of Auckland's Strong AI Lab and contributes to the PyTorch ONNX exporter, ONNX Runtime GenAI and Microsoft PyRIT.