torch.export and AOTInductor

What it is

torch.export is the supported way to capture a model into a stable, serializable graph. AOTInductor (torch._inductor.aoti_compile_and_package and friends) is the ahead-of-time compilation path that takes an exported program and produces a self-contained .so (or .pt2 package) plus a small C runtime that runs without Python.

Together they replace TorchScript as the recommended deployment path for PyTorch 2.x.

import torch
exported = torch.export.export(model, args=(x,), dynamic_shapes={"x": {0: torch.export.Dim("batch")}})
torch.export.save(exported, "model.pt2")

# Later, load and AOT-compile:
ep = torch.export.load("model.pt2")
so_path = torch._inductor.aot_compile(ep.module(), args=(x,))

Pipeline

graph LR
    Model[nn.Module] -->|torch.export.export| Tr[Dynamo + ProxyTorchDispatchMode]
    Tr -->|joint trace| EP[ExportedProgram]
    EP -->|serde schema| Pt2[.pt2 file]
    EP -->|aot_compile| Inductor[Inductor]
    Inductor -->|generate code| Cpp[C++/Triton kernels]
    Cpp -->|build| So[Compiled .so]
    So -->|aoti_runtime| Run[Python or C++ runtime]

ExportedProgram

ExportedProgram (in torch/export/exported_program.py) is the captured artefact:

• graph_module — an FX GraphModule.
• graph_signature — describes inputs (user inputs, parameters, buffers, constants) and outputs (user outputs, mutated buffers, gradients).
• range_constraints — symbolic-shape constraints accumulated during tracing.
• module_call_graph — preserves the original nn.Module call structure for post-processing.
• state_dict — parameter and buffer values.
• verifier — runtime invariants the program must satisfy.

ExportedProgram is not a nn.Module. To run it eagerly, call .module() to get a callable wrapper.

What export does differently from compile

• No graph breaks. export is fullgraph=True by definition; if the model can't be fully traced, export raises.
• Joint forward+backward optional. By default export traces only forward; for training-time export use torch.export.export_for_training.
• Stable schema. The captured graph is validated against torch/_export/serde/schema.py. Schema versions are forward/backward compatible.
• Pre-decompositions optional. export produces a graph in the core ATen op set or a smaller set as configured.
• Predictable shapes. All dynamic dims are explicit Dim annotations; the constraint system is deterministic.

Dynamic shapes

dynamic_shapes={"x": {0: Dim("batch", min=1, max=128)}, "y": {1: Dim.STATIC}} tells the tracer exactly which dims may vary. Dynamic dims become SymInts flowing through the graph and are constrained by the shape environment. After export, the range_constraints field summarizes the inferred constraints.

Decompositions

Export exposes a decomp_table (a dict of OpOverload -> decomposition function) that controls which ATen ops are reduced to simpler ones. The default decomposes everything down to core ATen — a stable subset of ~250 ops. Backends like ExecuTorch can provide their own table to widen or narrow the supported op set.

AOTInductor

torch._inductor.aot_compile (and the higher-level aoti_compile_and_package) feeds an ExportedProgram to Inductor and asks it to emit a deployable artefact:

• A C++ wrapper (in torch/csrc/inductor/aoti_runtime/) that exposes a run function.
• One or more Triton or C++ kernel sources, compiled against PyTorch headers.
• A weight blob.

The whole thing builds into a single .so (or .pt2 package containing the .so plus serialized weights). The runtime is small enough to ship with non-Python applications.

The aoti_runtime C API is in torch/csrc/inductor/aoti_runtime/. There is also a thin Python loader (torch._inductor.aoti_load_package) for easy use from Python.

Use cases

• ExecuTorch. Mobile/edge inference; ExecuTorch consumes ExportedProgram and produces a .pte flat file plus a C++ runtime.
• Server inference at Meta. AOTInductor .sos are deployed without Python on internal serving stacks.
• ONNX export. The Dynamo-onnx exporter consumes ExportedProgram. See ONNX.
• Safety-critical deployment. Stable schema + verifier means upgrades don't silently change behaviour.

Where it lives

Path	Contents
torch/export/	Public API (export, Dim, save, load)
torch/_export/	Implementation
torch/_export/serde/schema.py	The serialization schema
torch/_export/passes/	Lowering passes
torch/_inductor/aot_compile.py	AOTInductor entry
torch/csrc/inductor/aoti_runtime/	AOTInductor C runtime
torch/_export/db/	Examples database used for testing

Where to read next

• Systems / AOT Autograd — the same machinery underneath.
• Systems / Inductor — the codegen.
• Systems / FX — IR.
• Systems / Serialization — how .pt2 files work.