torch.compiler.substitute_in_graph

torch.compiler.substitute_in_graph(original_fn, *, can_constant_fold_through=False, skip_signature_check=False)[源代码]

为一个函数注册一个垫片处理程序，通常这个函数是C扩展中的C函数。当在图中内联原始函数时，使用该处理程序来替代原始函数。

注意

polyfill 处理程序仅在内联原始函数时使用。直接调用原始函数时则不使用它。在急切模式下，装饰后的函数会调用高性能的 C 函数，而不是 polyfill 处理程序。

当内联原始函数时，会使用 polyfill 处理函数来替换原始函数。polyfill 处理函数应与原始函数具有相同的签名和行为。

参数

• original_fn (callable) – 需要注册多填充处理程序的原始函数，通常是一个C函数。

• can_constant_fold_through (bool, 可选) – 是否可以通过 polyfill 处理程序进行常量折叠。具体来说，如果 polyfill 处理程序是纯函数且其参数为常量，则在编译期间可以对 polyfill 处理程序的结果进行常量折叠。默认值为 False。

• skip_signature_check (bool, 可选) – 是否跳过原始函数和 polyfill 处理程序之间的签名检查。默认为 False。

返回值

一个装饰器，用于为原始函数注册 polyfill 处理程序。

返回类型

Callable [[_F], _F]

示例：

>>> 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