Torch Export

torch.export

Capture a PyTorch nn.Module into an ExportedProgram -- a fully traced ATen-dialect graph with no Python runtime dependency. This is the standard entry point for deploying PyTorch models to any runtime (ExecuTorch, TensorRT, ONNX, custom backends).

Basic Export

from torch.export import export

exported = export(model.train(False), (example_input,))

Dynamic Shapes

By default every dimension is static. Use dynamic_shapes to allow variable sizes at runtime:

from torch.export import Dim, export

batch = Dim("batch", min=1, max=32)
seq_len = Dim("seq_len", min=1, max=2048)

exported = export(
    model.train(False),
    (example_input,),
    dynamic_shapes={"x": {0: batch, 1: seq_len}},
)

Rules:

• Keys match kwarg names or positional index
• Values map dimension indices to Dim objects
• Same Dim object across inputs asserts equal size
• Dim.AUTO for best-effort dynamic marking

The ExportedProgram

print(exported.graph)            # torch.fx.Graph -- the traced DAG
gm = exported.graph_module       # torch.fx.GraphModule wrapping the graph
out = exported.module()(x)       # run eagerly for correctness testing

IR Levels

Most deployment backends use Core ATen IR.

Serialization

torch.export.save(exported, "model.pt2")
loaded = torch.export.load("model.pt2")

Control Flow

• Static: branches on Python primitives are resolved at trace time
• Shape-dependent: branches on dynamic shapes emit guards (must be provable)
• Data-dependent: use torch._check to assert one branch, or torch.cond for both

Higher-order operators: torch.cond, torch.while_loop, torch.map, torch.scan, torch.associative_scan

Debugging Failures

• Draft export: torch.export.draft_export(model, args) -- always produces a graph, reports all issues
• Verbose logs: TORCH_LOGS="+dynamo,+export" python -m your_export_module 2>&1 | tlparse
• Incremental: test static shapes first, then add one dynamic dimension at a time
• Auto-fix shape constraints: refine_dynamic_shapes_from_suggested_fixes(str(e), ds)

Reference

Load guides/torch-export.md when the user needs a deep dive on tracing internals, symbolic shapes, guards, control flow operators, common failure patterns, or making untraceable code traceable.