PyTorch 2.x torch.export 完整 API 参考：从导出、序列化到动态形状与反扁平化

torch.export 是 PyTorch 编译器栈（TorchDynamo + AOTAutograd + Inductor）的入口级 API，它以"提前编译（AOT）"方式把任意 nn.Module 与示例输入捕获为一张只描述 Tensor 计算、可跨环境序列化、可对后续输入保持形状约束有效性的中间表示。本篇以仓库中 api_reference.md 所列 API 为主线，结合 torch/export 源码，逐一讲解 export、ExportedProgram、动态形状（Dim/dims/ShapesCollection/AdditionalInputs）、save/load、draft_export、unflatten 与 register_dataclass 的签名、参数语义、底层行为与可运行的代码示例，帮助读者把模型导出、部署与二次开发链路打通。

一、torch.export 模块全景

torch/export 是 torch.export 用户可见 API 的公共门面，其 __all__ 完整导出了本文要介绍的公开符号（见 torch/export/init.py）：

__all__ = [
    "AdditionalInputs", "Constraint", "CustomDecompTable", "default_decompositions",
    "Dim", "dims", "draft_export", "export", "ExportBackwardSignature",
    "ExportedProgram", "ExportGraphSignature", "FlatArgsAdapter", "load",
    "ModuleCallEntry", "ModuleCallSignature", "register_dataclass", "save",
    "ShapesCollection", "unflatten", "UnflattenedModule",
]

按功能可以把这些 API 分为五组，与 api_reference.md 中的 autodoc 指令一一对应：

分组	API	对应源码模块
核心导出入口	export、ExportedProgram、default_decompositions	torch/export/init.py、torch/export/exported_program.py
动态形状	Dim、dims、Constraint、ShapesCollection、AdditionalInputs、refine_dynamic_shapes_from_suggested_fixes	torch/export/dynamic_shapes.py
序列化	save、load、package_pt2、load_pt2	torch/export/init.py、torch/export/pt2_archive/_package.py
诊断与重构	draft_export、unflatten、UnflattenedModule、FlatArgsAdapter	torch/export/_draft_export.py、torch/export/unflatten.py
类型扩展	register_dataclass、custom_ops、custom_obj、decomp_utils、graph_signature	torch/export/init.py、torch/export/custom_ops.py、torch/export/graph_signature.py

其中 decomp_utils、graph_signature、pt2_archive.constants 等模块主要由内部机制使用，用户日常接触较少，本文重点覆盖前四组。

二、核心入口：torch.export.export

2.1 函数签名与参数说明

export 定义于 torch/export/init.py：

def export(
    mod: torch.nn.Module,
    args: tuple[Any, ...],
    kwargs: Mapping[str, Any] | None = None,
    *,
    dynamic_shapes: Any = None,
    strict: bool = False,
    preserve_module_call_signature: tuple[str, ...] = (),
    prefer_deferred_runtime_asserts_over_guards: bool = False,
) -> ExportedProgram:

各参数语义如下（源码 docstring 原义，见 torch/export/init.py）：

• mod：被追踪的模块，只会追踪其 forward 方法。
• args / kwargs：示例位置输入与关键字输入，用于驱动追踪并记录输入树结构（TreeSpec）。
• dynamic_shapes：动态形状规格，支持三种形式（详见第三节）：
  1. 从输入参数名到规格的 dict；
  2. 按原函数签名顺序排列的 tuple / list；
  3. 更新的结构化 API torch.fx.experimental.dynamic_spec.ShapesSpec / ParamsSpec（与 torch.compile 的 dynamic_shapes= 同一套规格，推荐用于后续开发）。
• strict：默认 False，走 Python 运行时追踪（pre-dispatch），仍会校验形状安全等关键假设；设为 True 时走 TorchDynamo 追踪以保证图的一致性（soundness），但 Dynamo 对 Python 特性覆盖有限，可能报更多错误。注意该开关不影响最终 IR 形态与序列化结果。
• preserve_module_call_signature：需要保留原始调用约定的子模块路径列表，这些元数据会在调用 torch.export.unflatten 时用于恢复子模块的原始调用方式。
• prefer_deferred_runtime_asserts_over_guards：偏好将形状约束生成为延迟运行时断言（deferred runtime assert）而非捕获期 guard。

2.2 导出产物与 Soundness 保证

export 的返回值是 ExportedProgram。源码 docstring 明确了三条核心保证（见 torch/export/init.py）：

1. 图只包含功能化后的 ATen 算子集合（以及用户指定的自定义算子）；
2. 消除了所有 Python 控制流与数据结构（少数例外）；
3. 记录了证明上述规范化与消除对"未来输入"成立所需的全部形状约束。

关于 soundness：追踪过程中 export 会记录用户程序与底层算子内核做出的形状假设，只有当这些假设对后续输入成立时，ExportedProgram 才有效。静态形状假设自动校验；动态形状假设必须通过 Dim 与 dynamic_shapes 显式声明。假设无法校验时会抛致命错误，错误信息会附带建议修复（suggested fixes），例如把 Dim("dim0_x") 改为 Dim("dim0_x", max=5)，可将其原样复制进代码重新导出（参见 torch/export/init.py）。

2.3 内部实现要点

export 的入口实现做了四件事（torch/export/init.py）：

1. 校验 mod 必须是 nn.Module，且不支持 torch.jit.ScriptModule（提示可用 TS2EPConverter 转换）；
2. 解析 mod.forward 上可能携带的 @dynamic_spec(...) 装饰器，将其中的 ShapesSpec 用作 dynamic_shapes（两者同时给出会报错）；
3. 委托 torch/export/_trace.py 中的 _export 完成实际追踪，pre_dispatch=True；
4. 捕获异常时，对 GuardOnDataDependentSymNode、UnsupportedOperatorException、UserError、ConstraintViolationError 等已知错误类型追加提示，建议用户改用 draft_export() 获取更完整的错误报告。

2.4 可接受输入/输出类型

export 支持三类输入/输出（torch/export/init.py）：

• 原始类型：torch.Tensor、int、float、bool、str；
• 先经 register_dataclass 注册的 dataclass；
• 由 dict、list、tuple、namedtuple、OrderedDict 构成的（可嵌套）数据结构。

最简导出示例（来自 save 的 docstring）：

import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

ep = torch.export.export(MyModule(), (torch.randn(5),))

三、动态形状：Dim 与 dynamic_shapes

3.1 Dim：声明动态维度的核心类

Dim 定义于 torch/export/dynamic_shapes.py，其 docstring 明确了两套用法：

（1）Dim Hints（自动动态形状）——Dim.AUTO、Dim.DYNAMIC、Dim.STATIC，门槛最低，用户只声明维度是动态、静态或交由编译器推断：

class Foo(nn.Module):
    def forward(self, x, y):
        assert x.shape[0] == 4
        assert y.shape[0] >= 16
        return x @ y

x = torch.randn(4, 8)
y = torch.randn(8, 16)
dynamic_shapes = {
    "x": {0: Dim.AUTO, 1: Dim.AUTO},
    "y": {0: Dim.AUTO, 1: Dim.AUTO},
}
ep = torch.export.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)

此时若全部换成 Dim.DYNAMIC 会报错，因为模型本身要求 x.shape[0] 是静态的。Hints 也支持给边界：Dim.AUTO(min=16, max=32)、Dim.DYNAMIC(max=64)，编译器会在该范围内推断其余约束；若有效范围完全落在用户指定范围之外会抛异常（见 torch/export/dynamic_shapes.py）。

底层实现中，三种 Hint 对应 _DimHintType 枚举（AUTO/STATIC/DYNAMIC，见 dynamic_shapes.py），在解析时分别映射为 maybe_mark_dynamic、mark_static、mark_dynamic 并生成 _RelaxedConstraint（dynamic_shapes.py）。

（2）Named Dims（命名维度）——Dim("name", min=..., max=...)，更严格，编译器推断出的约束与用户声明不一致会直接报错：

s0 = Dim("s0")
s1 = Dim("s1", min=16)
dynamic_shapes = {
    "x": {0: 4, 1: s0},
    "y": {0: s0, 1: s1},
}
ep = torch.export.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)

构造约束：min 缺省为 0，max 缺省为 int_oo（无穷），且必须满足 max > min；名称必须是合法 Python 标识符（dynamic_shapes.py）。

3.2 维度间关系：派生 Dim

Named Dims 还支持声明维度之间的单变量线性关系，例如：

s0 = Dim("s0")
s1 = 3 * s0 + 4

这是通过 Dim 重载的 +、-、* 运算符（仅限整数系数）实现的，运算结果是一个 _DerivedDim（dynamic_shapes.py）。_DerivedDim 恒可写成 Ax + B（A 为正整数），保证单调递增，从而简化取值范围推导、等式判定等元理论（dynamic_shapes.py）。关系式（例如 (x.shape[0] + y.shape[1]) % 4 == 0）可能被编译为图中的运行时断言节点，在运行时输入不满足约束时触发断言。

dims 工具函数用于批量创建命名 Dim：

def dims(*names: str, min: int | None = None, max: int | None = None) -> tuple[Dim, ...]:
    return tuple(Dim(name, min=min, max=max) for name in names)

即 s0, s1 = dims("s0", "s1", min=2)（见 dynamic_shapes.py）。

3.3 dynamic_shapes 的三种书写形式

dynamic_shapes 参数整体支持（torch/export/init.py）：

• 按参数名的 dict：{"x": {0: s0}, "y": [s0, s1]}；
• 按顺序的 tuple/list：[{0: s0}, [s0, s1]]，关键字参数须按原函数签名顺序；
• 嵌套递归：dict/list/tuple 中的张量参数可用 dict（维度索引 → Dim/None）或 tuple/list（每维一个 Dim 或 None）递归描述；不需要包含静态维度，包含时须映射为 None（注意源码同时提示：用 None 表示静态维度已废弃，请改用 Dim.STATIC，见 dynamic_shapes.py）。

此外 dynamic_shapes 还接受 torch.fx.experimental.dynamic_spec.ShapesSpec / ParamsSpec（见 torch/export/init.py），其特性包括：仅支持 unbacked SymInt（u 符号，永不特化）、维度可以是符号表达式（如 TensorSpec([batch * 2, 3])）、支持符号间关系假设（assumptions=[batch % 2 == 0]）且无静默特化——导出图对每条假设都保证有效，否则导出失败。示例：

batch = ShapeVar("batch", min=2, max=128)
ep = torch.export.export(
    mod,
    (torch.randn(8, 3), torch.randn(16, 3)),
    dynamic_shapes=ShapesSpec(
        params=ParamsSpec({
            "x": TensorSpec([batch, 3]),
            "y": TensorSpec([batch * 2, 3]),  # 派生表达式
        }),
        assumptions=[batch % 2 == 0],
    ),
    strict=True,
)

注意：draft_export 目前不支持 ShapesSpec/ParamsSpec（见下文第五节的实现限制）。

3.4 形状规格校验与约束生成

• _check_dynamic_shapes（dynamic_shapes.py）负责校验：dict 顶层键必须与输入参数名一致（否则给出详细提示，包括"是否忘了给单个 dict 输入包一层 list/tuple"这类常见错误）；张量形状规格的长度必须与实际张量维数一致；同一命名 Dim 在不同位置出现时 min/max 必须一致；非张量输入只能映射 None。
• _process_dynamic_shapes（dynamic_shapes.py）把规格解析为约束列表：命名 Dim → _Constraint（含 StrictMinMaxConstraint 值域）；派生 Dim → _DerivedConstraint（其根可能是真实输入维度或"幽灵根" _PhantomRoot，例如形状 2*dim 与 dim+1 通过一个不直接出现在输入形状中的中间符号关联）；Hints → _RelaxedConstraint。
• 形状校验失败时的建议修复可以用 refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes) 自动吸收进新规格（dynamic_shapes.py）。它解析错误消息中 Suggested fixes: 之后的 dim = Dim('dim', min=3, max=6)、dim = 4、dy = dx + 1 等条目，返回可直接再次传给 export 的新 dynamic_shapes：

try:
    ep = export(mod, args, dynamic_shapes=dynamic_shapes)
except torch._dynamo.exc.UserError as exc:
    new_shapes = refine_dynamic_shapes_from_suggested_fixes(exc.msg, dynamic_shapes)
    ep = export(mod, args, dynamic_shapes=new_shapes)

3.5 面向复杂输入的辅助构建器

当 args 是嵌套结构、按结构书写 dynamic_shapes 很繁琐时，有两个专门工具：

ShapesCollection（dynamic_shapes.py）：按张量对象直接赋值形状规格，内部用 id(tensor) 作索引，__setitem__ 支持 torch.Tensor 与包装整数的 _IntWrapper；同一张量重复赋值且规格不一致会报错。dynamic_shapes(m, args, kwargs) 生成与 args 结构一致的规格树；若有赋值过的张量没出现在 args 中会抛 ValueError（提示可能因拷贝或未注册 pytree 的类导致）。示例：

args = {"x": tensor_x, "others": [tensor_y, tensor_z]}

dim = torch.export.Dim("dim")
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[tensor_y] = {0: dim * 2}
# 等价于 dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [{0: dim * 2}, None]}

torch.export.export(mod, args, dynamic_shapes=dynamic_shapes)

AdditionalInputs（dynamic_shapes.py）：面向部署工程师——不关心模型内部逻辑，只提供一批有代表性的测试/剖析输入，由工具自动推断哪些形状应动态。与原始示例输入形状不同的维度视为动态，相同则视为静态；verify(ep) 会逐条校验导出程序对每个附加输入都有效（复用 _check_input_constraints_for_module），保证用它们追踪会得到同一张图：

args0, kwargs0 = ..., ...      # 导出用示例输入
dynamic_shapes = torch.export.AdditionalInputs()
dynamic_shapes.add(args1, kwargs1)   # 追加代表性输入
...
torch.export.export(mod, args0, kwargs0, dynamic_shapes=dynamic_shapes)

推断逻辑对整数值也做了处理：int 在多次输入中取值不同才标记 Dim.DYNAMIC；非 int 值不同则直接报错（dynamic_shapes.py）。

四、导出产物：ExportedProgram

4.1 类结构与核心属性

ExportedProgram（torch/export/exported_program.py）"打包"了一次 export 的全部产物：一个表示 Tensor 计算的 torch.fx.Graph、包含所有被提升（lift）参数与缓冲区张量值的 state_dict，以及各种元数据。类字段（均带类型注解的 docstring）包括：

• _graph_module：底层 torch.fx.GraphModule；
• _graph_signature：ExportGraphSignature，图的输入/输出规格；
• _state_dict：原模块参数与缓冲区值；
• _range_constraints：动态形状的符号约束（sympy.Symbol → ValueRanges）；
• _module_call_graph：模块层级与调用签名（list[ModuleCallEntry]）；
• _example_inputs：导出时使用的 (args, kwargs)；
• _constants：图用到的常量（含非持久缓冲区、常量张量、自定义对象）；
• _verifiers：用于校验导出程序的校验器类列表。

构造过程会剥离图的 codegen 相关设置（保证是平铺图）、执行 _common_getitem_elimination_pass 公共优化、并在构造最后调用 self.validate() 校验（exported_program.py）。

4.2 常用属性与方法

• 只读属性：graph_module、graph、graph_signature、state_dict、range_constraints、module_call_graph、example_inputs、call_spec、verifier/verifiers、constants。所有 setter 均直接抛 RuntimeError，防止外部破坏导出图的不变量。
• parameters() / named_parameters() / buffers() / named_buffers()：按 graph_signature 中记录的参数、缓冲区名遍历原模块的参数与缓冲区（非持久缓冲区从 constants 中取，见 exported_program.py）。
• module(check_guards=True)：返回自包含的 GraphModule，所有参数/缓冲区被内联回去（内部走 _unlift_exported_program_lifted_states）。check_guards=True（默认）时生成 _guards_fn 子模块并在占位符后插入输入 guard 检查；False 时部分检查由前向 pre-hook 完成。返回的模块 train()/eval() 被替换为抛 NotImplementedError 的实现（exported_program.py）。
• __call__ 被禁用：直接调用 ExportedProgram(...) 会抛 RuntimeError，必须改用 ep.module()（exported_program.py）。
• run_decompositions(decomp_table=None, decompose_custom_triton_ops=False)：对导出图运行算子分解并返回新 ExportedProgram。默认执行 Core ATen 分解（default_decompositions()）使算子落入 Core ATen 算子集；decomp_table={} 表示不分解任何算子；也可在 default_decompositions() 基础上覆盖个别算子（见 exported_program.py）。CustomDecompTable 通过 materialize() 展开成实际表。

default_decompositions() 定义于 exported_program.py，返回 CustomDecompTable，是 run_decompositions 默认行为的依据。

4.3 图签名：ExportGraphSignature 与相关类型

graph_signature 模块（torch/export/graph_signature.py）定义了 ExportGraphSignature 及一组规格类型：ArgumentSpec、InputKind/InputSpec、OutputKind/OutputSpec、ConstantArgument、TensorArgument、SymIntArgument、SymFloatArgument、SymBoolArgument、CustomObjArgument、TokenArgument，以及 ExportBackwardSignature（torch/export/init.py 导出）。

InputKind 区分用户输入、参数、缓冲区、常量张量、自定义对象等类别；ExportGraphSignature 上提供 parameters、buffers、non_persistent_buffers、inputs_to_buffers、buffers_to_mutate、user_inputs_to_mutate 等映射，ExportedProgram 的 named_parameters/named_buffers/_graph_module_flat_inputs 均依赖这些映射实现状态查找与输入拼接。调用约定为"先参数、后缓冲区、再用户输入"（见 exported_program.py 中的 additional_inputs + flat_args）。

五、序列化：save / load 与 pt2_archive

5.1 torch.export.save

签名（torch/export/init.py）：

def save(
    ep: ExportedProgram,
    f: FileLike,
    *,
    extra_files: dict[str, Any] | None = None,
    opset_version: dict[str, int] | None = None,
    pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
) -> None

• f：实现 write/flush 的文件类对象，或文件名（str | os.PathLike | IO[bytes]）；
• extra_files：{文件名: 内容} 映射，会作为额外文件一并存入归档；
• opset_version：opset 名到版本的映射；
• pickle_protocol：覆盖默认 pickle 协议（默认 2）；
• 实现上委托 pt2_archive._package.package_pt2，以 {"model": ep} 作为导出程序集合打包（torch/export/init.py）。

ep = torch.export.export(MyModule(), (torch.randn(5),))

# 存到文件
torch.export.save(ep, "exported_program.pt2")
# 存到 io.BytesIO 缓冲
buffer = io.BytesIO()
torch.export.save(ep, buffer)
# 附带额外文件
extra_files = {"foo.txt": "bar"}
torch.export.save(ep, "exported_program.pt2", extra_files=extra_files)

5.2 torch.export.load

签名（torch/export/init.py）：

def load(
    f: FileLike,
    *,
    extra_files: dict[str, Any] | None = None,
    expected_opset_version: dict[str, int] | None = None,
) -> ExportedProgram

• f：与 save 相同的文件类对象或文件名；
• extra_files：给出期望读取的额外文件名（值会被实际内容替换）；
• expected_opset_version：期望的 opset 版本映射，用于版本校验。

ep = torch.export.load("exported_program.pt2")

with open("exported_program.pt2", "rb") as f:
    buffer = io.BytesIO(f.read())
buffer.seek(0)
ep = torch.export.load(buffer)

extra_files = {"foo.txt": ""}   # 值将被数据替换
ep = torch.export.load("exported_program.pt2", extra_files=extra_files)
print(extra_files["foo.txt"])

加载流程（torch/export/init.py）优先走新格式 load_pt2，失败时回退到旧 zip 格式：读取 version 与 schema 版本比对，再反序列化 serialized_exported_program.json、serialized_state_dict.pt、serialized_constants.pt、serialized_example_inputs.pt，并对 extra_files 前缀条目做解包。torch/export/pt2_archive/constants.py 定义归档内部的文件名常量。

安全警告（源码 docstring 原义）：torch.export.load 底层使用 pickle，切勿加载不可信来源的数据。同时两个 API 均标注"正在积极开发中，旧文件在新版本中可能不可用"。

仓库测试 test/export/test_serialize.py 覆盖了 save → 缓冲 → load → 再执行的完整往返（例如 L191-L193 等用例），是序列化链路的直接验证。

六、诊断利器：draft_export

draft_export（torch/export/init.py）与 export 签名一致，但设计目标相反：只要可能，就尽量产出一个 ExportedProgram，同时生成一份列出所有潜在 soundness 问题的报告。这在排查导出失败时非常有用——export 抛错时其异常信息也会主动提示"把 export() 换成 draft_export() 查看更完整的信息"（torch/export/init.py）。

实现要点（torch/export/_draft_export.py）：

• 通过 torch._functorch.config.patch 开启 fake_tensor_propagate_real_tensors=True 与 generate_fake_kernels_from_real_mismatches=True，让没有 fake kernel 的算子也能继续追踪并记录 profile；
• 捕获 ConstraintViolationError 等异常，尝试按建议修复动态形状后重试，把成功路径与失败列表写入 DraftExportReport；
• 报告挂在 ep._report 上：report.successful() 判断是否完全成功，print(ep._report) 在 Python 中查看错误，失败时会提示可生成 HTML 报告页；
• 对追踪中发现的、未注册 fake kernel 的算子（report.op_profiles），在 unsafe_generate_fake_kernels 上下文内生成临时 fake kernel 完成导出。

限制：draft_export 不支持新的 ShapesSpec/ParamsSpec 动态形状 API，传入会抛 NotImplementedError（torch/export/init.py）。

七、模块层级重构：unflatten

unflatten(module: ExportedProgram, flat_args_adapter: FlatArgsAdapter | None = None) -> UnflattenedModule（torch/export/unflatten.py）把扁平图恢复成与原 eager 模块相同的模块层级，便于接入期望"模块树"而非"单张平图"的推理系统。

• UnflattenedModule（torch/export/unflatten.py）：一个 torch.nn.Module 子类，其 process_forward_inputs 被 torch._dynamo.disable(..., recursive=True) 包裹，避免把输入预处理逻辑纳入追踪（torch/export/unflatten.py）。
• FlatArgsAdapter（torch/export/unflatten.py）：抽象基类，用于在输入 TreeSpec 与导出模块不匹配时适配扁平化参数。
• 重要限制（docstring 原义）：反扁平化后模块的 args/kwargs 不一定与 eager 模块一致，直接做模块替换（如 self.submod = new_mod）不一定可行；需要替换子模块时，应在 export 时设置 preserve_module_call_signature 保留原调用约定。该选项的元数据以 ModuleCallEntry / ModuleCallSignature（输入输出规格 + in/out TreeSpec + 前向参数名，见 torch/export/exported_program.py）形式记录在 module_call_graph 中，供 unflatten 使用。
• 实现还包含 _inplace_buffer_and_input_mutations 等图变换：把 functionalization 产生的"缓冲区/输入变更"从"输入-输出贯通"形式改写为 aten.copy_ 原地节点，使反扁平化模块的行为贴近原始 eager 代码（torch/export/unflatten.py）。

八、类型扩展与其余公开 API

8.1 register_dataclass

register_dataclass(cls, *, serialized_type_name=None)（torch/export/init.py）把 dataclass 注册为 export 的合法输入/输出类型，实现上直接调用 pytree.register_dataclass。serialized_type_name 用于希望序列化包含该 dataclass 的 pytree TreeSpec 的场景。示例（docstring 原义）：

import torch
from dataclasses import dataclass

@dataclass
class InputDataClass:
    feature: torch.Tensor
    bias: int

@dataclass
class OutputDataClass:
    res: torch.Tensor

torch.export.register_dataclass(InputDataClass)
torch.export.register_dataclass(OutputDataClass)

class Mod(torch.nn.Module):
    def forward(self, x: InputDataClass) -> OutputDataClass:
        res = x.feature + x.bias
        return OutputDataClass(res=res)

ep = torch.export.export(Mod(), (InputDataClass(torch.ones(2, 2), 1),))
print(ep)

8.2 custom_ops 与 custom_obj

• torch/export/custom_ops.py：在 FRAGMENT 库 "export" 下定义 access_subclass_inner_tensor(Tensor, str) -> Tensor 等内部辅助算子，并把自定义 autograd 函数按字符串名导入调用（_call_custom_autograd_function_in_pre_dispatch），用于 pre-dispatch 追踪中保持自定义算子语义。
• torch/export/custom_obj.py：管理导出图中出现的自定义 Python 对象，序列化时作为常量（constants/CustomObjArgument）一并打包，保证 load 后程序仍可运行。

8.3 其余模块

• decomp_utils（torch/export/decomp_utils.py）：CustomDecompTable 是 run_decompositions 中分解表的容器（支持 None/空 dict/自定义 dict 三种语义）；
• torch.export.passes（torch/export/passes）：导出后置 pass 的公共命名空间，ExportedProgram 构造与反扁平化中会用到的若干 pass 也在此注册；
• pt2_archive 与 pt2_archive.constants（torch/export/pt2_archive）：package_pt2 / load_pt2 是 save/load 的底层实现，constants 定义归档文件布局常量，仅在深度调试序列化格式时需要关注。

九、端到端工作流小结

把以上 API 串成一条典型的生产链路：

import torch

# 1) 定义模型与示例输入
class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(8, 4)
    def forward(self, x):
        return self.fc(x).relu()

model = Net().eval()
x = torch.randn(2, 8)

# 2) 声明动态维度并导出
b = torch.export.Dim("b", min=1, max=64)
ep = torch.export.export(model, (x,), dynamic_shapes={"x": {0: b, 1: None}})

# 3) 校验：查看图签名与形状约束
print(ep.graph)
print(ep.range_constraints)

# 4) 序列化 / 反序列化
torch.export.save(ep, "net.pt2")
ep2 = torch.export.load("net.pt2")

# 5) 反扁平化恢复模块层级（如需）
unflat = torch.export.unflatten(ep2)

# 6) 推理：获取可调用 GraphModule
gm = ep2.module()
out = gm(torch.randn(3, 8))   # 动态 batch 维度可变化

至此，读者应已掌握 torch.export 全套公开 API：用 export + Dim 体系做 AOT 捕获与动态形状声明，用 ExportedProgram 访问图/签名/状态与执行分解，用 save/load 做跨进程、跨版本的模型交付，用 draft_export 定位导出失败根因，用 unflatten 对接模块树形态的推理框架，并用 register_dataclass 扩展自定义数据类型。进一步深入时，可继续阅读仓库中的 torch/export/_trace.py（追踪内核）、torch/export/_unlift.py（状态内联）以及测试目录 test/export（尤其 test_export.py 与 test_serialize.py）中的用例来印证各 API 的实际行为。