PyTorch torch.export IR 规范详解：ExportedProgram、Graph、Node 与元数据契约

torch.export 是 PyTorch 用于将模型以 AOT（Ahead-of-Time）方式捕获为编译期中间表示（IR）的核心入口。本文以仓库内 docs/source/user_guide/torch_compiler/export/ir_spec.md 为骨架，系统讲解 Export IR 的 ExportedProgram、Graph、Node、SymInt、FakeTensor 与 Pytree-able 类型体系，并结合 torch/export/exported_program.py、torch/export/graph_signature.py、torch/export/__init__.py 等源码，帮助你从"会调用 torch.export.export()"进阶到"能读懂、能校验、能为其编写后端"的层次。

Export IR 是什么

Export IR 是一种面向编译器的图中间表示，设计目标与 MLIR、TorchScript 类似，但专门用于表达 PyTorch 程序的语义。它把计算表达为一条精简的操作列表，仅对控制流等动态行为提供有限支持（例如通过 cond、map 等高阶算子表达条件与循环）。

它的产生流程是：

1. 前端通过 trace-specializing（跟踪特化）机制可靠地（soundly）捕获一个 PyTorch 程序；
2. 得到的 Export IR 图随后可被后端优化与执行。

当前仓库中，这一过程由 torch.export.export() 完成。其 docstring 明确给出了三条关键承诺：生成的功能化 ATen 算子（以及用户自定义算子）是归一化的、Python 控制流与数据结构（在特定例外下）被消除、并记录了一组形状约束以保证这种归一化与控制流消除对未来输入是可靠的（soundness guarantee）。

与 torch.fx 的关系：更强的严格性

Export IR 是实现在 torch.fx.Graph 之上的——"所有 Export IR 图同时也是合法的 FX 图"，若按标准 FX 语义解释，Export IR 可被可靠地解释。一个直接推论是：导出的图可以经由标准 FX codegen 转换回合法的 Python 程序。

因此文档刻意把篇幅集中在 Export IR 与 FX 的差异点（更严格的约束） 上，而不是重复 FX 已有内容。理解这一前提，是读懂本文后续所有"差异"条款的基础。

ExportedProgram：IR 的顶层容器

Export IR 的顶层构造是 torch.export.ExportedProgram 类，它把模型的计算图与模型消费的参数/权重打包在一起。仓库中类定义见 torch/export/exported_program.py#L1070-L1153，其 docstring 概括为："包含 torch.fx.Graph 表示的 Tensor 计算、包含所有被提升参数与 buffer 数值的 state_dict，以及各类元数据"，并且可以像原始可调用对象一样被调用。

关键属性

属性	类型	说明
graph_module	torch.fx.GraphModule	存放拍平后的计算图；图可直接通过 ExportedProgram.graph 访问
graph_signature	torch.export.ExportGraphSignature	记录图中使用与改写的参数、buffer 名称。参数与 buffer 不再作为图模块的属性存在，而是被提升为图的输入；signature 负责跟踪这些参数的附加信息
state_dict	Dict[str, Union[torch.Tensor, torch.nn.Parameter]]	存放参数与 buffer 数值的数据结构
range_constraints	Dict[sympy.Symbol, RangeConstraint]	对带有数据依赖行为导出的程序，节点元数据含符号形状（形如 s0、i0）；该属性把符号映射到其上下界

从源码看，ExportedProgram 的构造函数在 exported_program.py#L1112-L1153 中依次完成：剥离 graph 上的 codegen 并创建扁平图模块、执行 _common_getitem_elimination_pass 公共 getitem 消除、保存 signature/state_dict/range_constraints/module_call_graph、默认装配 Verifier 校验器，并以 self.validate() 收尾——校验始终是构造函数的最后一步。同时，graph_module、graph、graph_signature、state_dict 等属性均为只读（setter 直接抛 RuntimeError），保证导出的程序不会被随意改写。

graph_signature：参数与 buffer 的"户口本"

ExportGraphSignature 定义于 torch/export/graph_signature.py#L165-L275。它建模了 Export Graph 的输入/输出签名，核心不变量是：

• Export Graph 是功能化的：图内不通过 getattr 节点访问参数、buffer 等"状态"；export 保证参数、buffer、常量张量被提升为图的输入；
• 对 buffer 的改写也不直接出现在图内，而是把改写后的 buffer 值建模为图额外的输出。

输入与输出的排序规则为：

Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs]
Outputs = [*mutated_inputs, *flattened_user_outputs]

InputKind（PARAMETER、BUFFER、CONSTANT_TENSOR、USER_INPUT、CUSTOM_OBJ、TOKEN 等）与 OutputKind（USER_OUTPUT、BUFFER_MUTATION、PARAMETER_MUTATION、USER_INPUT_MUTATION、GRADIENT_TO_PARAMETER 等）分别枚举输入、输出的类别，parameters、buffers、non_persistent_buffers、lifted_tensor_constants、user_inputs、user_outputs 等属性则从 input_specs/output_specs 中按类别抽取对应名称集合（见 graph_signature.py#L277-L374）。

例如一个同时含参数、两个 buffer 并在 forward 中原地改写 buffer 的模块，其 graph_signature 会呈现为：

# inputs
p_my_parameter: PARAMETER target='my_parameter'
b_my_buffer1: BUFFER target='my_buffer1' persistent=True
b_my_buffer2: BUFFER target='my_buffer2' persistent=True
x1: USER_INPUT
x2: USER_INPUT

# outputs
add_2: BUFFER_MUTATION target='my_buffer2'
add_1: USER_OUTPUT

该示例完整代码与图对照见 graph_signature.py#L183-L271，其中 run_decompositions() 会把非功能化图（含 aten.add_.Tensor 原地算子）转换为功能化图（buffer 改写变为额外输出 add_2）。这正是理解"Export IR 为什么能放心交给后端执行"的关键：图内不再有隐藏状态。

Graph：DAG 与节点构成

一个 Export IR Graph 是 DAG（有向无环图）形式的 PyTorch 程序：每个节点代表一次计算/操作，边由节点间的引用构成。其 schema 极简：

class Graph:
  nodes: List[Node]

实践中 Export IR 的图即 torch.fx.Graph Python 类。一个合法的 Export IR 图必须包含：

• 0 个或多个 op 类型为 placeholder 的节点；
• 0 个或多个 op 类型为 call_function 的节点；
• 恰好 1 个 op 类型为 output 的节点。

推论：最小的合法 Graph 只有一个节点，即 nodes 永不为空。

定义：placeholder 节点集合代表 Graph/GraphModule 的输入；output 节点代表其输出。

一个最简可运行示例（对应文档原始示例）：

import torch
from torch import nn

class MyModule(nn.Module):

    def forward(self, x, y):
      return x + y

example_args = (torch.randn(1), torch.randn(1))
mod = torch.export.export(MyModule(), example_args)
print(mod.graph)

输出（Graph 的文本表示，每行一个节点）：

graph():
  %x : [num_users=1] = placeholder[target=x]
  %y : [num_users=1] = placeholder[target=y]
  %add : [num_users=1] = call_functiontarget=torch.ops.aten.add.Tensor, kwargs = {})
  return (add,)

Node：节点 Schema 与 FX 文本格式

Node 表示一次具体计算/操作，Python 层面用 torch.fx.Node 类表示。节点间的边通过 Node.args 属性中对其他节点的直接引用实现；借助 FX 机制可表达计算图所需的操作符调用、占位符（输入）、条件分支与循环等。

Node 的 schema：

class Node:
  name: str # name of node
  op_name: str  # type of operation

  # interpretation of the fields below depends on op_name
  target: [str|Callable]
  args: List[object]
  kwargs: Dict[str, object]
  meta: Dict[str, object]

FX 文本格式

如上例所示，文本格式的每行遵循：

%<name>:[...] = <op_name>target=<target>), kwargs = {"keyword": arg5})

该格式以紧凑形式囊括了 Node 类的全部字段（meta 除外）。具体地：

• <name>：即 node.name；
• <op_name>：即 node.op 字段，必须是 call_function、placeholder、get_attr、output 之一；
• <target>：即 node.target，其含义取决于 op_name；
• args：即 node.args 元组中的内容；若列表中某值是 torch.fx.Node，会特别以 % 前缀标识。

例如一个 add 算子的调用显示为：

%add1 = call_functiontarget = torch.op.aten.add.Tensor, kwargs = {})

其中 %x、%y 是两个名字为 x、y 的 Node。值得注意：字符串 torch.op.aten.add.Tensor 表示的是 target 字段中实际存储的可调用对象，而非仅仅是它的字符串名。

文本格式的最后一行 return [add] 对应 op_name = output 的节点，表示返回这一元素。

call_function：算子调用节点

call_function 节点表示对某个算子的调用。

两个关键定义：

• Functional（功能化）：一个可调用对象满足以下全部条件即为 functional——非改写（不修改输入的值，对张量而言包括元数据与数据）且无副作用（不改写外部可见的状态，如模块参数值）；
• Operator（算子）：具有预定义 schema 的 functional 可调用对象，典型如功能化 ATen 算子。

FX 表示：

%name = call_functiontarget = operator, kwargs = {})

与 vanilla FX call_function 的差异：

1. FX 图中 call_function 可指向任意可调用对象；Export IR 则限定为 ATen 算子的选定子集、自定义算子与控制流算子；
2. Export IR 中常量参数会被内嵌进图内；
3. FX 图中 get_attr 节点可读取图模块中存储的任何属性；Export IR 中它被限定为只能读取子模块，因为所有参数/buffer 都以输入形式传入图模块。

call_function 的元数据契约

Node.meta 是挂在每个 FX 节点上的 dict，但 FX 规范并不规定其中会有哪些元数据。Export IR 提供了更强的契约：所有 call_function 节点保证恰好包含以下字段：

• node.meta["stack_trace"]：字符串，引用原始 Python 源码的调用栈，例如：

  File "my_module.py", line 19, in forward
  return x + dummy_helper(y)
  File "helper_utility.py", line 89, in dummy_helper
  return y + 1
  

• node.meta["val"]：描述该操作运行后的输出，类型可为 <symint>、<FakeTensor>、List[Union[FakeTensor, SymInt]] 或 None；
• node.meta["nn_module_stack"]：该节点来源的 torch.nn.Module 的"调用栈"。例如 addmm 算子出自 nn.Sequential 内的 nn.Linear，则表现为：

  {'self_linear': ('self.linear', <class 'torch.nn.Linear'>), 'self_sequential': ('self.sequential', <class 'torch.nn.Sequential'>)}
  

• node.meta["source_fn_stack"]：分解（decomposition）之前该节点来自的 torch 函数或叶子 torch.nn.Module 类。例如 addmm 来自 nn.Linear 模块调用，则 source_fn 为 torch.nn.Linear；来自 torch.nn.functional.linear 函数调用则包含 torch.nn.functional.Linear。

这些元数据使后端能够在编译期获得源码定位、输出类型形状、模块层级与原始算子来源——这正是 Export IR 能支撑高质量错误报告与算子级优化的基础。

placeholder：图输入节点

placeholder 表示图的输入，语义与 FX 完全一致。占位节点必须是图 nodes 列表中最靠前的 N 个节点（N 可以为 0）。

FX 表示：

%name = placeholdertarget = name)

target 字段是输入名称的字符串；args 若非空，则大小为 1，表示该输入的默认值。

元数据：placeholder 节点同样有 meta['val']（与 call_function 一致），此时 val 表示图期望为该输入参数接收的输入 shape/dtype。

output：图出口节点

output 调用对应函数中的 return 语句，因此它终止当前图。图中有且仅有一个 output 节点，且它永远是图的最后一个节点。

FX 表示：

output[](args = (%something, …))

语义与 torch.fx 完全一致，args 表示要返回的节点。output 节点的元数据与 call_function 节点相同。

get_attr：子模块读取节点

get_attr 节点表示从封装它的 torch.fx.GraphModule 中读取一个子模块。与 torch.fx.symbolic_trace 得到的 vanilla FX 图不同——后者用 get_attr 从顶层 GraphModule 读取参数、buffer 等属性——Export IR 中参数与 buffer 均作为输入传入图模块，并存储在顶层 ExportedProgram 中，因此 get_attr 在这里仅用于读取子模块（典型场景是控制流分支的嵌套子图）。

FX 表示：

%name = get_attrtarget = name)

示例（使用 functorch.experimental.control_flow.cond 构造条件分支）：

from functorch.experimental.control_flow import cond

def true_fn(x):
    return x.sin()

def false_fn(x):
    return x.cos()

def f(x, y):
    return cond(y, true_fn, false_fn, [x])

对应图：

graph():
    %x_1 : [num_users=1] = placeholder[target=x_1]
    %y_1 : [num_users=1] = placeholder[target=y_1]
    %true_graph_0 : [num_users=1] = get_attr[target=true_graph_0]
    %false_graph_0 : [num_users=1] = get_attr[target=false_graph_0]
    %conditional : [num_users=1] = call_functiontarget=torch.ops.higher_order.cond, kwargs = {})
    return conditional

其中 %true_graph_0 = get_attr[target=true_graph_0] 读取包含 sin 算子的子模块 true_graph_0，而 torch.ops.higher_order.cond 即为文档所提及的控制流算子之一。

引用类型：SymInt、FakeTensor 与 Pytree-able 类型

SymInt

SymInt 要么是字面整数，要么是表示某个整数的符号（Python 中由 sympy.Symbol 类表示）。当 SymInt 是符号时，它描述的是一个编译期未知、仅运行时可知的整型变量——这正是动态 shape 在 Export IR 中的载体。

FakeTensor

FakeTensor 是包含张量元数据的对象，可视为：

class FakeTensor:
  size: List[SymInt]
  dtype: torch.dtype
  device: torch.device
  dim_order: List[int]  # 尚不存在

• size 是整数或 SymInt 的列表：若含 SymInt 表示该张量具有动态形状；若全是整数则假定张量具有该精确静态形状。张量元数据的秩（rank）永不是动态的；
• dtype 表示该节点输出的数据类型；Export IR 中不存在隐式类型提升（no implicit type promotions）；
• FakeTensor 中没有 strides。

换言之，node.meta['val'] 与 node.target 算子的返回类型一一对应：

• 算子返回 Tensor → val 为描述该张量的 FakeTensor；
• 算子返回 n 元组张量 → val 为 n 元 FakeTensor；
• 算子返回编译期已知的 int/float/scalar → val 为 None；
• 算子返回编译期未知的 int/float/scalar → val 为 SymInt。

例如：aten::add 返回 Tensor，其 spec 即带 dtype 与 size 的 FakeTensor；aten::sym_size 返回整数，其 val 是 SymInt（值仅运行时可得）；max_pool2d_with_indexes 返回 (Tensor, Tensor) 元组，spec 就是 2 元 FakeTensor，分别描述返回值的两个元素。

Python 代码示例：

def add_one(x):
  return torch.ops.aten(x, 1)

对应图：

graph():
  %ph_0 : [#users=1] = placeholder[target=ph_0]
  %add_tensor : [#users=1] = call_functiontarget=torch.ops.aten.add.Tensor, kwargs = {})
  return [add_tensor]

FakeTensor：

FakeTensor(dtype=torch.int, size=[2,], device=CPU)

Pytree-able 类型

一个类型是 Pytree-able 的，当且仅当它是叶子类型或包含其他 Pytree-able 类型的容器类型（pytree 概念与 JAX 的 pytrees 文档一致）。

叶子类型（leaf type）：

类型	定义
Tensor	torch.Tensor
Scalar	Python 中任意数值类型，包括整型、浮点型与零维张量
int	Python int（C++ 中绑定为 int64_t）
float	Python float（C++ 中绑定为 double）
bool	Python bool
str	Python 字符串
ScalarType	torch.dtype
Layout	torch.layout
MemoryFormat	torch.memory_format
Device	torch.device

容器类型（container type）：

类型	定义
Tuple	Python tuple
List	Python list
Dict	以 Scalar 为键的 Python dict
NamedTuple	Python namedtuple
Dataclass	必须通过 register_dataclass 注册（实现见 torch/export/init.py）
Custom class	通过 _register_pytree_node 定义的任意自定义类（实现见 torch/utils/_pytree.py）

该类型体系决定了 export 边界上可合法传递的数据结构：用户在 torch.export.export() 中传入的 args/kwargs、模型返回的结果，最终都会被按 pytree 规则拍平为叶子节点，再映射为图的 placeholder 与 output。

总结

Export IR 的核心设计可以浓缩为三个关键词：

1. 功能化——参数、buffer 被提升为输入，buffer 改写被建模为输出，图中不存在隐藏状态（见 graph_signature.py#L166-L176）；
2. 严格元数据契约——stack_trace、val、nn_module_stack、source_fn_stack 四字段为每个 call_function 节点背书，编译期即可获得完整语义；
3. 受限的算子集合与显式类型——call_function 仅限 ATen 算子子集、自定义算子与控制流算子，常量内嵌，类型经 SymInt/FakeTensor 显式刻画，无隐式提升。

无论你是为 Export IR 编写自定义后端，还是调试 torch.export.export() 的输出图，ir_spec.md 与其对应源码（exported_program.py、graph_signature.py、torch/export/init.py）都是最权威的第一手资料。