What export refuses

Chapter 7 frames export as the thing that refuses the compromise dynamo makes. The frame is right and it hides a fork: there are two ways to export, and they fail differently because they are built out of different machinery.

Strict mode, the default in torch 2.2.2, traces with dynamo. It analyses bytecode, so it knows which Python line offended and says so. Non-strict mode runs your forward on the actual Python interpreter with fake tensors flowing through it, so anything that is not a tensor operation simply executes and leaves no trace in the graph.

Put a print inside a forward and both behaviours show up in one run. Strict refuses with Unsupported: call_function BuiltinVariable(print), and it names the two arguments it saw, a string constant and a tuple, because it read the call out of the bytecode rather than executing it. Non-strict accepts, and the graph that comes back has four nodes with no print among them.

The line tracing (3,) lands on your terminal twice, so the forward ran twice inside the one export call. The shape in it is not a guess. tuple(x.shape) was evaluated on the fake tensor the interpreter was carrying, and a fake tensor holds an exact shape, dtype and device with no data behind them. The shape was there to read. The values were not, which is what stops the second module.

import torch

class Logged(torch.nn.Module):
    def forward(self, x):
        w = x.tanh()
        print("tracing", tuple(x.shape))
        return w * 3

class Guarded(torch.nn.Module):
    def forward(self, x):
        if (x < 0).any():
            return x.abs()
        return x.tanh()

for name, mod in [("logged", Logged()), ("guarded", Guarded())]:
    for strict in (True, False):
        try:
            ep = torch.export.export(mod, (torch.ones(3),), strict=strict)
            print(f"{name} strict={strict}: accepted, {len(ep.graph.nodes)} nodes")
        except Exception as e:
            print(f"{name} strict={strict}: {type(e).__name__}: {str(e).splitlines()[0][:110]}")

The output below is verbatim, tracing lines included, in the order they were printed. Read the two guarded lines together. Both modes refuse the data-dependent if, and the wording of the refusal tells you which tracer produced it.

Strict mode raises a UserError naming the feature and pointing at functorch.experimental.control_flow.cond, the operator that would let the branch live inside the graph as a node with two subgraphs. Non-strict raises DataDependentOutputException naming aten._local_scalar_dense.default, which is the op that pulls a Python number out of a tensor. (x < 0).any() produces a one-element bool tensor and the if needs a real bool out of it, which means reading the data a fake tensor does not have.

Same refusal, two different vantage points on it, and the practical value is diagnostic. A UserError quoting your source line means dynamo tripped. A fake-tensor exception naming an ATen op means you were in non-strict mode and the interpreter got as far as executing before the shape machinery stopped it.

logged strict=True: Unsupported: call_function BuiltinVariable(print) [ConstantVariable(str), TupleVariable()] {}
tracing (3,)
tracing (3,)
logged strict=False: accepted, 4 nodes
guarded strict=True: UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to ex
guarded strict=False: DataDependentOutputException: aten._local_scalar_dense.default

Now the case that goes the other way, and it is the one worth remembering because it inverts the usual story. The guard-or-break drill on the gym's pytorch floor has a .item() scenario among its eight, and its measured verdict is a graph break. .item() has to hand back a real Python number, and at capture time there is no such number, so capture stops at that line and Python runs it.

Export takes the same construct. Subtract a mean pulled out with .item() and the graph that comes back is one placeholder, three ATen calls and a return. The middle of the three is aten._local_scalar_dense.default, the scalar extraction written down as an operation instead of appearing as a gap in the capture. Nothing broke because there was nothing to break: export has no fallback path to Python, so it either represents a thing in the graph or refuses.

The line between the two cases is whether the value gets used to decide control flow. A scalar pulled out and fed back into arithmetic is a node. The same scalar compared in an if is a branch, and a branch needs both sides in the graph, which is what cond is for and what the error message asks you to write.

import torch

class Center(torch.nn.Module):
    def forward(self, x):
        return x - x.mean().item()

print(torch.export.export(Center(), (torch.ones(4),)).graph)

# graph():
#     %l_x_ : [num_users=2] = placeholder[target=l_x_]
#     %mean : [num_users=1] = call_function[target=torch.ops.aten.mean.default](args = (%l_x_,), kwargs = {})
#     %_local_scalar_dense : [num_users=1] = call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%mean,), kwargs = {})
#     %sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%l_x_, %_local_scalar_dense), kwargs = {})
#     return (sub,)

Chapter 6 proves that a Python side effect inside a compiled function keeps firing on every call, and contrasts that with jax's trace-once model where it fires exactly once. Export lands on the jax side of that line, and it is worth measuring rather than assuming.

Append the input's shape to a module-level list, export, and the list holds one entry, (3,): the trace ran the forward once, on the example input you handed it. Empty the list, replay the exported program three times, and it stays empty. The graph is a placeholder, one tanh, and the return, with no record that a Python list was ever touched.

That is the contract you are buying. An ExportedProgram is a description of tensor computation and nothing else, which is what lets it run in a process with no Python interpreter. Anything your forward did that was not tensor math happened once, during export, in your process, and is gone.

import torch

seen = []

class Recorder(torch.nn.Module):
    def forward(self, x):
        seen.append(tuple(x.shape))
        return x.tanh()

ep = torch.export.export(Recorder(), (torch.ones(3),))
print("seen after export:", seen)
seen.clear()
for _ in range(3):
    ep.module()(torch.ones(3))
print("seen after three replays:", seen)
print(ep.graph)

# seen after export: [(3,)]
# seen after three replays: []
# graph():
#     %l_x_ : [num_users=1] = placeholder[target=l_x_]
#     %tanh : [num_users=1] = call_function[target=torch.ops.aten.tanh.default](args = (%l_x_,), kwargs = {})
#     return (tanh,)

One graph for the whole program means one graph for every input shape you intend to serve, so export makes you say which dimensions vary. The module below is chapter 7's Sin, the same two ops it exports there, taken back to the export call with one argument added. A Dim with a name and a range turns a concrete size into a symbol, and the placeholder's shape comes back as (s0,) instead of (6,).

The range is not decoration. ep.range_constraints records it, and the exported module enforces both ends at replay: a size of 3 or of 9 against a Dim("b", min=4, max=8) raises Input l_x_.shape[0] is outside of specified dynamic range [4, 8], while 4 and 8 run. Compare that to the dynamo layer, where an unexpected shape is a guard miss and a silent recompile.

The two failure modes are the same disagreement chapter 6 draws between guarded capture and tracing once, met a second time at the artifact boundary. Under compile, a shape you did not anticipate costs you a compile. Under export, it costs you an exception, which is what you want from something running in a serving stack with no compiler attached.

import torch
from torch.export import Dim

class Sin(torch.nn.Module):
    def forward(self, x):
        return torch.sin(x).sum()

b = Dim("b", min=4, max=8)
ep = torch.export.export(Sin(), (torch.randn(6),), dynamic_shapes={"x": {0: b}})
print("range_constraints:", ep.range_constraints)
print("placeholder shape:", tuple(list(ep.graph.nodes)[0].meta["val"].shape))
for n in (3, 4, 8, 9):
    try:
        print(n, "->", round(ep.module()(torch.ones(n)).item(), 6))
    except Exception as e:
        print(n, "->", type(e).__name__ + ":", str(e).splitlines()[0][:70])

# range_constraints: {s0: ValueRanges(lower=4, upper=8, is_bool=False)}
# placeholder shape: (s0,)
# 3 -> RuntimeError: Input l_x_.shape[0] is outside of specified dynamic range [4, 8]
# 4 -> 3.365884
# 8 -> 6.731767
# 9 -> RuntimeError: Input l_x_.shape[0] is outside of specified dynamic range [4, 8]

The graph is the part everyone looks at, and it is not the part that makes an ExportedProgram portable. Parameters do not live inside the graph as attributes; they are lifted to placeholders, and a separate graph_signature records which placeholder was which parameter. That is why the printed graph of a two-parameter linear layer starts with three inputs when the module takes one.

Alongside the signature there is a state_dict holding the actual weights under their original names, so nothing about the mapping is lost. torch.export.save writes the whole thing to a stream, and that stream is a zip archive with four entries: the serialized program, the state dict, any constants, and a version file. Reload it in another process and the module replays with no reference to the class that produced it.

This is the object chapter 10 hands to the bridges. What each bridge does with it after that, and the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → it becomes on the other side, is that chapter's material and the torch_xla lessons underneath it. From here, the thing worth carrying is what the artifact is: a graph, a signature that maps its inputs back to names, and the weights.

import io
import zipfile
import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)
    def forward(self, x):
        return torch.relu(self.fc(x)).sum()

torch.manual_seed(0)
ep = torch.export.export(Net(), (torch.randn(2, 3),))
for spec in ep.graph_signature.input_specs:
    print(spec.kind, spec.arg.name, "->", spec.target)
print("state_dict keys:", list(ep.state_dict))
buf = io.BytesIO()
torch.export.save(ep, buf)
buf.seek(0)
print("archive entries:", zipfile.ZipFile(buf).namelist())
buf.seek(0)
print("reloaded, replayed:", round(torch.export.load(buf).module()(torch.ones(2, 3)).item(), 6))

# InputKind.PARAMETER arg0_1 -> fc.weight
# InputKind.PARAMETER arg1_1 -> fc.bias
# InputKind.USER_INPUT l_x_ -> None
# state_dict keys: ['fc.weight', 'fc.bias']
# archive entries: ['serialized_exported_program.json', 'serialized_state_dict.json', 'serialized_constants.json', 'version']
# reloaded, replayed: 0.314205