Finding ONNX exporter bugs with differential testing: what worked and what fooled me
The PyTorch ONNX exporter translates each ATen op through a function in onnxscript's torchlib. Those functions have unit tests, but the tests mostly feed ordinary random floats. I wanted to know what happens at the edges: NaN, infinity, signed zero, integer and boolean inputs, float16. So I wrote a small differential test that runs each op in PyTorch eager mode and through torch.onnx.export(dynamo=True) plus ONNX Runtime, and compares the two.
Over three rounds and roughly 300 op-and-argument variants, it found several real bugs, four of which I fixed in open pull requests. It also produced four kinds of false alarm, which are the more useful part of this post.
The harness
Each case is a tiny nn.Module that calls one op. The harness runs it eagerly, exports it, runs the ONNX model on ONNX Runtime, and compares shape, dtype and values. Floats are compared with a tolerance and equal_nan=True; integers and booleans must match exactly.
class M(torch.nn.Module):
def __init__(self, fn): super().__init__(); self.fn = fn
def forward(self, *args): return self.fn(*args)
expected = M(fn)(*inputs)
program = torch.onnx.export(M(fn), inputs, dynamo=True)
actual = program(*inputs) # runs on ONNX Runtime
compare(actual, expected) # shape, dtype, values
The inputs matter more than the ops. Each op ran on a float tensor containing NaN, ±inf, ±0 and tiny values, on int32, int64 and uint8 tensors, on booleans, on an empty tensor, and in float16, bfloat16 and float64. Most bugs live where a PyTorch op accepts a dtype or a special value that the matching ONNX operator does not.
What it found
| Finding | Effect | Status |
|---|---|---|
ceil, floor, round, trunc, isnan, isfinite on integer inputs | Invalid ONNX model | My fix: onnxscript#3178 |
max(x, dim) and min(x, dim) on bool | Invalid ONNX model | My fix: onnxscript#3179 |
heaviside(NaN, v) | Returns 1; PyTorch returns 0 | My fix: onnxscript#3180 |
softmax/log_softmax with dtype= | Casts the output instead of the input: zeros or an invalid model for integer input | My fix: onnxscript#3196 |
atan2 with signed zeros, infinities and NaN | Wrong angles | Already being fixed in onnxscript#3152 |
cross_entropy(label_smoothing=...) | Wrong loss | Already being fixed in onnxscript#3069 |
scatter_reduce(reduce="mean") | Overwrites instead of averaging | Already being fixed in onnxscript#3094 |
special.erfcx for large inputs | 0 or NaN | Already being fixed in onnxscript#3038 |
The same idea at model level, exporting more than 40 Hugging Face models and comparing their outputs, found one more: Swin and BEiT cannot be exported with SDPA attention. I wrote that one up in a separate post.
Half of the real findings were already claimed by someone else. I learned to search open and closed pull requests before writing any code. Once, I wrote a full atan2 fix with tests before finding that another contributor had opened the same fix days earlier.
Four ways it fooled me
- Constants get folded before ONNX Runtime sees them. In the second round I built the inputs inside
forward. The exporter's optimizer then folded the whole computation at export time, using the ONNX reference implementation, so I was partly testing that evaluator, not ONNX Runtime.logsumexpover a row of-infcame out as NaN that way. ONNX Runtime itself returns-inf, and the reference implementation had already been fixed upstream (onnx#8441). In the third round every input was a graph input. - Tolerance artefacts.
logit(sigmoid(10 * x), eps=1e-6)showed a 0.05 difference. With exact inputs, the export matches PyTorch:sigmoidsaturates close to 1 in float32, andlogitamplifies the last bits. Re-test every mismatch on hand-picked inputs before believing it. - The runtime can hide a wrong graph. For float16 input,
softmax(x, dtype=torch.float32)exported asSoftmax(float16) -> Cast, yet the numbers were right on CPU. ONNX Runtime's CPU provider has no float16Softmaxkernel and upcasts on its own, so the bug only shows on providers that do have one. The regression test therefore checks the graph structure, with the optimizer turned off. - Not every error is a bug. Most failures in the third round were missing ops (
digamma,linalg_inv,bessel_i0,masked_selectand others), which are feature requests. Some were my own mistakes: a random generator insideforwardcannot be exported at all.
Takeaways
- A short harness and a good list of edge-case inputs find real bugs quickly. Exotic inputs matter more than exotic ops.
- Feed the values as graph inputs, re-check every hit by hand, and look at the exported graph when the numbers look too good.
- Search for existing issues and pull requests before writing a fix.
I used Claude (Anthropic) to help build the harness, 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.