torch.compiler.substitute_in_graph

torch.compiler.substitute_in_graph(original_fn, *, can_constant_fold_through=False, skip_signature_check=False)[source][source]

注册一个用于替换原始函数的 polyfill 处理器，通常是一个来自 C 扩展的 C 函数，当在图中内联原始函数时使用。

注意

当内联原始函数时，才会使用 polyfill 处理器。如果直接调用原始函数，则不会使用它。在急切模式下，装饰的函数调用性能更好的 C 函数，而不是 polyfill 处理器。

polyfill 处理器是一个函数，当内联原始函数时，将代替原始函数被调用。polyfill 处理器应具有与原始函数相同的签名和相同的行为。

参数:

• original_fn (可调用对象) – 要为其注册 polyfill 处理器的原始函数，通常是一个 C 函数。

• can_constant_fold_through (布尔值，可选) – polyfill 处理器是否可以被常量折叠。也就是说，如果 polyfill 处理器是一个纯函数，并且其参数是常量，那么在编译期间，polyfill 处理器的结果可以被常量折叠。默认值为 False 。

• 跳过签名检查（布尔值，可选）- 是否跳过原始函数和 polyfill 处理程序之间的签名检查。默认为 False 。

返回值:

用于注册原始函数的 polyfill 处理程序的装饰器。

返回类型:

Callable[[Callable[[~_P], _R]], Callable[[~_P], _R]]

示例：

>>> import operator
>>> operator.indexOf([1, 2, 3, 4, 5], 3)
2
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
... # xdoctest: +SKIP("Long tracebacks")
Traceback (most recent call last):
...
torch._dynamo.exc.Unsupported: ...

>>> @torch.compiler.substitute_in_graph(operator.indexOf)
... def indexOf(a, b, /):
...     for i, item in enumerate(a):
...         if item is b or item == b:
...             return i
...     raise ValueError("sequence.index(x): x not in sequence")
>>>
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
2