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+电感器）
没有装饰器	跟踪内部	跟踪内部
允许在图中	不可见可调用对象	内部追踪
自定义操作	不可见调用	不可见调用

allow_in_graph() 的一个常见用例是作为编译器前端的一个逃生门：如果你知道该函数在编译堆栈的下游组件（AOTAutograd 和 Inductor）方面是有效的，但 Dynamo 存在一个错误，阻止它正确地符号性地内省该函数（或者如果你的代码是 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 内部使用的所有张量必须直接传递给 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() 的单个图。