ONNX

Active contributors: titaiwangms, xadupre, justinchuby

Purpose

torch.onnx is PyTorch's exporter to the ONNX interchange format. Two exporter implementations coexist:

1. Legacy / TorchScript-based (torch.onnx.export(..., dynamo=False)) — traces through TorchScript and walks JIT IR ops to ONNX nodes via per-op symbolic functions in torch/onnx/symbolic_*.py.
2. Dynamo-based (torch.onnx.export(..., dynamo=True) or torch.onnx.dynamo_export) — uses Dynamo to capture an FX graph, decomposes it through the ATen op set, and lowers each ATen op via the onnxscript library. The recommended path going forward.

Directory layout

Path	Contents
torch/onnx/	Python package
torch/onnx/__init__.py	Public surface (export, dynamo_export, register_custom_op_symbolic)
torch/onnx/utils.py	Legacy export entry point
torch/onnx/_internal/	Internal infrastructure for the dynamo exporter
torch/onnx/_internal/exporter/	Dynamo exporter
torch/onnx/_internal/diagnostics/	SARIF diagnostics for export failures
torch/onnx/symbolic_opset*.py	Per-opset (1, 9, 11, 13, 17, 18, …) symbolic functions for legacy export
torch/onnx/symbolic_helper.py	Common helpers used by symbolic_opset files
torch/onnx/symbolic_caffe2.py	Legacy Caffe2 export (removed-in-progress)
torch/csrc/jit/passes/onnx*	C++ JIT passes that prepare graphs for legacy export
test/onnx/	ONNX export test suite

Key abstractions

Type	File	Purpose
torch.onnx.export	torch/onnx/__init__.py	User-facing entry point
OnnxExporterError	torch/onnx/errors.py	The exception users see
_internal.exporter.Exporter	torch/onnx/_internal/exporter/	Dynamo-path orchestration
OpsetRegistry	torch/onnx/_internal/registration.py	Symbolic-function registry per opset
SymbolicContext	torch/onnx/symbolic_helper.py	The handle each symbolic function receives

How it works

Legacy path

graph LR
    User[Python model + args] -->|jit.trace or jit.script| JIT[TorchScript IR]
    JIT -->|ONNX prep passes| JIT2[Cleaned IR]
    JIT2 -->|per-op symbolic functions| ONNX[ONNX graph]
    ONNX -->|onnx.save| File[.onnx file]

The legacy exporter:

1. Traces the model through TorchScript to produce JIT IR.
2. Runs a series of cleanup passes (torch/csrc/jit/passes/onnx*): inlining, peephole, constant folding, scalar-type promotion, eval-mode insertion.
3. For each ATen op in the graph, looks up a symbolic function — usually in torch/onnx/symbolic_opset<n>.py for the targeted opset — and calls it with a context to emit one or more ONNX nodes.

Symbolic functions are the main contributor surface for the legacy exporter; PyTorch ships hundreds of them.

Dynamo path

graph LR
    User[Python model + args] -->|Dynamo| FX[FX graph]
    FX -->|decompositions| FX2[ATen-only FX graph]
    FX2 -->|per-op onnxscript translators| ONNX[ONNX graph]
    ONNX --> File[.onnx file]

The Dynamo exporter:

1. Uses Dynamo (or torch.export) to produce an ATen-level FX graph.
2. Applies decompositions to reduce the graph to the core ATen op set.
3. Lowers each remaining op via onnxscript's registered translators.
4. Optionally uses ONNX shape inference + ONNX-runtime to validate before serialization.

The Dynamo path inherits all of torch.compile's tracing improvements: dynamic shapes, custom ops, higher-order ops.

Custom op support

For custom ops, users register a custom symbolic:

@torch.onnx.symbolic_helper.parse_args("v")
def my_op(g, x):
    return g.op("MyDomain::MyOp", x)

torch.onnx.register_custom_op_symbolic("my_lib::my_op", my_op, 17)

In the Dynamo path, custom ops are registered with onnxscript instead.

Integration points

• JIT. The legacy path relies on JIT IR and JIT passes. See JIT.
• Dynamo / FX. The new path consumes FX graphs from Dynamo. See Dynamo, FX.
• Quantization. PT2E quantized models can be exported via the Dynamo path; quantized-FX export through the legacy path is in torch/onnx/symbolic_opset10.py.
• onnxscript and onnxruntime are runtime deps installed alongside PyTorch when ONNX export is used.

Entry points for modification

• New op support in legacy export → add a symbolic function in the appropriate symbolic_opset<n>.py.
• New op support in Dynamo export → add or update an entry in onnxscript's torchlib (out of tree) or in the in-tree decomposition tables.
• New custom op registration → register_custom_op_symbolic.

Key source files

File	Purpose
torch/onnx/__init__.py	Public surface
torch/onnx/utils.py	Legacy export
torch/onnx/symbolic_opset17.py	Recent-opset symbolic functions
torch/onnx/_internal/exporter/	Dynamo-path implementation
torch/csrc/jit/passes/onnx.cpp	JIT-side ONNX prep