torch.compiler.allow_in_graph

torch.compiler.allow_in_graph(fn)

告诉编译器前端 (Dynamo) 跳过对该函数的符号内省，并在遇到时直接将其写入图中。

如果您正在使用 torch.compile()（后端为 “inductor”（默认））或 torch.export.export()，并尝试在整个追踪过程中将某个 Python 函数视为黑盒，请勿使用此 API。请转而创建自定义算子（参见PyTorch 自定义算子着陆页）。

警告

如果您是典型的 torch.compile 用户（例如，您正在将 torch.compile 应用于模型以使其运行更快），您可能不希望使用此函数。allow_in_graph() 是一个危险功能，因为它跳过了负责执行安全检查（图中断、处理闭包等）的编译器前端 (Dynamo)。不正确的使用将导致难以调试的静默错误问题。

给定一个没有 allow_in_graph 装饰器的 Python 函数，torch.compile 的常规执行会追踪进入函数内部。allow_in_graph() 会改变这一点，使前端不追踪进入函数内部，但编译器后端仍会追踪通过该函数。将其与自定义算子进行比较，自定义算子在整个 torch.compile 堆栈中将函数视为黑盒。下表比较了这些机制。

机制	前端 (Dynamo)	后端 (AOTAutograd+Inductor)
无装饰器	追踪内部	追踪内部
allow_in_graph	不透明可调用对象	追踪内部
自定义算子	不透明可调用对象	不透明可调用对象

allow_in_graph() 的一个常见用例是作为编译器前端的逃生舱：如果您知道该函数对于编译堆栈的下游组件（AOTAutograd 和 Inductor）来说是有效的，但存在 Dynamo 的错误阻止它正确地符号内省该函数（或者如果您的代码是用 C/C++ 编写，因此无法使用 Dynamo 进行内省），那么可以使用 allow_in_graph() 装饰该函数以绕过 Dynamo。

我们要求 fn 遵循以下限制。不遵守将导致未定义行为。

• fn 的输入必须是 FX 图中可代理 (Proxy-able) 的类型。有效类型包括：Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?] Tuple[Tensor?, …]/Tuple[int?, …]/Tuple[float?, …]/torch.dtype/torch.device

• fn 的输出必须是 FX 图中可代理 (Proxy-able) 的类型（参见上一条）。

• fn 内部使用的所有 Tensor 必须直接作为输入传递给 fn（而不是作为捕获的变量）。

参数

fn – 表示要包含在图中的可调用对象。如果 fn 是可调用对象的列表或元组，则它会递归地将 allow_in_graph() 应用于每个函数，并返回包含修改后函数的新列表或元组。

示例

torch.compiler.allow_in_graph(my_custom_function)

@torch.compile(...)
def fn(x):
    x = torch.add(x, 1)
    x = my_custom_function(x)
    x = torch.add(x, 1)
    return x

fn(...)

将捕获一个包含 my_custom_function() 的单一图。