skip to content

What does torch.export.export() return, and how is it stricter than tracing?

level: middleimportance: should knowfreq 50%

answer

  1. a data structure, not a callable file
  2. weights lifted out, mutations removed
  3. no graph break allowed here
  4. dimensions specialize unless you declare them
  5. branches become torch.cond

basics

~20 s

torch.export.export() returns an ExportedProgram: one whole-graph ATen representation plus a signature describing which inputs are parameters, buffers and user arguments, and symbolic shape constraints. Unlike tracing, it refuses to silently specialize — unsupported dynamism raises instead.

solid answer

~50 s

`torch.export.export(mod, args)` produces an `ExportedProgram`, not a callable file: a single functionalized ATen-level FX graph, a `graph_signature` that separates parameters and buffers from real user inputs, the flattened state, and the symbolic-shape guards under which the graph is valid. Two things make it stricter than `torch.jit.trace`. First, it is whole-graph or nothing — there is no partial capture and no eager fallback, so anything the tracer cannot represent fails at export time rather than at 3 a.m. Second, dynamism is declared rather than guessed: you pass `dynamic_shapes` with `torch.export.Dim` to say which dimensions vary, and any other dimension is specialized to the example and recorded as a guard. Genuine data-dependent branching has to be written as `torch.cond` instead of a Python `if`. You serialize with `torch.export.save` to a `.pt2` file, and the same `ExportedProgram` is the front door to ahead-of-time Inductor compilation and to on-device lowering.

code

python · 16 lines
python
import torch

class M(torch.nn.Module):
    def forward(self, x):
        return x.relu().sum(dim=1)

ep = torch.export.export(
    M(),
    (torch.randn(4, 8),),
    dynamic_shapes={"x": {0: torch.export.Dim("batch", min=1, max=64)}},
)

print(type(ep))                      # ExportedProgram
print(ep.graph_signature.input_specs)  # params/buffers vs user inputs
print(ep.module()(torch.randn(9, 8)).shape)  # torch.Size([9])
torch.export.save(ep, "m.pt2")

go deeper

for a junior

Know that torch.export is the modern capture API and that it returns an ExportedProgram object which you save with torch.export.save. Being able to name it as the successor to TorchScript capture is enough here.

for a middle

Explain what is inside an ExportedProgram — an ATen-level functionalized graph, a signature separating weights from inputs, and shape guards — and why whole-graph capture with no fallback is a deliberate strictness rather than a limitation.

for a senior

Demonstrate that you have fought a real export: declaring dynamic_shapes with Dim, chasing the operator that specialized a dimension, and rewriting a data-dependent branch as torch.cond. Say what you verify after export and what the guards mean for the consumer's contract.

for a principal

Own the choice to standardize capture on torch.export so that quantization, on-device lowering and ONNX all share one front end, and the cost of that: model code must be written export-clean, which is a review standard and a training concern, not a per-release scramble.

## What the call gives you back `torch.export.export(module, args, kwargs=None, dynamic_shapes=None)` returns an `ExportedProgram`. It is a data structure, not a runtime and not a file. Its parts matter because interviewers ask what is actually inside: - **`graph_module`** — an FX `GraphModule` whose nodes are ATen operators. The graph is *functionalized*: mutations and aliasing have been rewritten into pure operations, so a consumer does not have to reason about in-place effects. - **`graph_signature`** — the mapping that says which graph inputs are parameters, which are buffers, which are user inputs, and which outputs are mutated buffers. Tracing gives you a callable and hides this; export exposes it, which is what makes downstream transformation (quantization, delegation, fusion) tractable. - **`state_dict` / constants** — the weights, lifted out of the graph rather than embedded, so they can be replaced or quantized without touching the graph. - **`range_constraints`** — the symbolic-shape guards. These record the conditions under which the captured graph is a valid program. You call `ep.module()` to get something you can run in Python for testing, and `torch.export.save(ep, "m.pt2")` / `torch.export.load` to move it as a file. ## Strictness, part one: whole graph or nothing `torch.compile` handles code it cannot capture by breaking the graph and running that piece in eager Python. `torch.export` cannot — the entire point is a self-contained artifact for a consumer that may have no Python at all. So there is no graph break, no fallback, and no partial success. Code with a `print` of a tensor value, a dependence on a Python object the tracer cannot represent, or a call into an unsupported library raises at export time. That is a feature. The `torch.jit.trace` failure mode is a program that runs and is wrong. The `torch.export` failure mode is a stack trace on your laptop while you still have the model open. ## Strictness, part two: shapes are declared, not observed By default every dimension of every example input is specialized to the value you passed. Export records that as a guard, and a consumer that violates it is out of contract. To make a dimension vary you say so: ``` dynamic_shapes={"x": {0: torch.export.Dim("batch", min=1, max=64)}} ``` The exporter then carries a symbolic size through the graph and will tell you where that symbol got constrained — often somewhere surprising, like a reshape with a hardcoded literal or a comparison against a length. Chasing those messages is most of the work of a first export, and the messages are the value: they are enumerating the places your model secretly assumed a fixed shape. ## Strictness, part three: control flow must be explicit A data-dependent `if` cannot survive into a static graph as a Python statement. Export makes you rewrite it as `torch.cond(pred, true_fn, false_fn, operands)`, which puts both branches into the graph as subgraphs and the predicate as a real tensor input. Both arms must agree on output shapes and dtypes. Compared with tracing, where the `if` silently disappears, this is more work up front and a correct artifact at the end. The `strict=` argument selects between two tracing frontends — a Dynamo-based one with heavier soundness checking and a Python-level one that is more permissive with unusual code. Know that the knob exists and that the two frontends fail on different things; check the release you are on for which is the default before quoting one in an interview. ## Where an ExportedProgram goes next Export is a capture step, not a deployment target. Downstream it feeds ahead-of-time compilation to a shared library, on-device lowering to a mobile runtime, PT2-style quantization passes that rewrite the graph, and the modern ONNX exporter, which runs export first and then translates the ATen graph into ONNX operators. That common front end is the reason `torch.export` is worth learning even if your eventual runtime is something else entirely: fix your model once, at the capture step, and every downstream path gets easier. ## Against the older paths `torch.jit.trace` records one execution — cheap, silent when wrong. `torch.jit.script` compiles a Python subset — control flow survives, but the compiler rejects a lot of ordinary Python and the whole TorchScript stack is in maintenance. `torch.export` produces a sound, shape-annotated, functionalized graph with an explicit contract, and is where PyTorch's export investment now goes. When you say "we export the model" in an interview on torch 2.13, this is what you should mean, and you should be able to say what breaks when you try.

  • Export fails saying a dimension you marked dynamic was constrained to a specific value. How do you proceed?
    Read the message for the operator that imposed the constraint — it names the line. Usually it is a reshape or view with a hardcoded literal, an indexing expression built from a Python int, or a comparison against a length. Rewrite that site to derive the size from the tensor (x.shape[0], -1 in a view) instead of a constant. If the constraint is genuinely required by the maths, relax the Dim's range or accept specialization and document the contract.
  • Why does export lift parameters and buffers out of the graph instead of embedding them as constants?
    Because downstream passes need to address them. Quantization has to replace weights with quantized tensors, delegation has to hand a subgraph plus its weights to a backend, and a serving stack may want to swap weights without re-exporting. The graph_signature makes each input's role explicit, so a transform can rewrite state without re-deriving which constants were weights — something a frozen TorchScript graph makes much harder.
  • Does an ExportedProgram run faster than the eager module?
    Not by itself. Export captures a graph; it does not generate kernels. Speed comes from what you run the graph through afterwards — an ahead-of-time compiler, a backend delegate, or a runtime with fused operators. Treating export as an optimization step is a common confusion; it is a portability and transformability step whose payoff is that a compiler can now see the whole program.

saying these in an interview costs you the question

  • torch.export is torch.jit.trace with a new name
  • An ExportedProgram is a servable runtime you deploy directly
  • Exported graphs accept any shape because export is dynamic
  • A Python if/else exports fine and picks the right branch at run time
  • Exporting a model makes it faster on its own

context

resolving links…