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 前端	后端（AOT Autograd + Inductor）
没有装饰器	内部追踪	内部追踪
允许在图表中显示	不透明调用对象	内部追踪
自定义操作	不透明调用对象	不透明调用对象

allow_in_graph() 的一个常见用例是作为编译前端的逃生机制：如果你知道某个函数在编译堆栈的下游组件（如 AOTAutograd 和 Inductor）中可以正常工作，但由于 Dynamo 中存在 bug 导致无法正确地对该函数进行符号化内省，或者你的代码是用 C/C++ 编写的且无法使用 Dynamo 进行内省，那么你可以用 allow_in_graph() 装饰该函数以绕过 Dynamo。

我们要求 fn 遵守以下限制。如果不遵守这些限制，可能会导致未定义的行为：

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

• fn 的输出必须是 FX 图中可以被代理的类型（参见前一条目）

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

参数

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

示例：

torch.compiler.allow_in_graph(my_custom_function)

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

fn(...)

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