Qiming Bao · Blog · 11 November 2026 · English / 中文

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

FindingEffectStatus
ceil, floor, round, trunc, isnan, isfinite on integer inputsInvalid ONNX modelMy fix: onnxscript#3178
max(x, dim) and min(x, dim) on boolInvalid ONNX modelMy fix: onnxscript#3179
heaviside(NaN, v)Returns 1; PyTorch returns 0My fix: onnxscript#3180
softmax/log_softmax with dtype=Casts the output instead of the input: zeros or an invalid model for integer inputMy fix: onnxscript#3196
atan2 with signed zeros, infinities and NaNWrong anglesAlready being fixed in onnxscript#3152
cross_entropy(label_smoothing=...)Wrong lossAlready being fixed in onnxscript#3069
scatter_reduce(reduce="mean")Overwrites instead of averagingAlready being fixed in onnxscript#3094
special.erfcx for large inputs0 or NaNAlready 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

  1. 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. logsumexp over a row of -inf came 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.
  2. Tolerance artefacts. logit(sigmoid(10 * x), eps=1e-6) showed a 0.05 difference. With exact inputs, the export matches PyTorch: sigmoid saturates close to 1 in float32, and logit amplifies the last bits. Re-test every mismatch on hand-picked inputs before believing it.
  3. The runtime can hide a wrong graph. For float16 input, softmax(x, dtype=torch.float32) exported as Softmax(float16) -> Cast, yet the numbers were right on CPU. ONNX Runtime's CPU provider has no float16 Softmax kernel 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.
  4. Not every error is a bug. Most failures in the third round were missing ops (digamma, linalg_inv, bessel_i0, masked_select and others), which are feature requests. Some were my own mistakes: a random generator inside forward cannot be exported at all.

Takeaways

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.