torch.compiler.substitute_in_graph¶
- torch.compiler.substitute_in_graph(original_fn, *, can_constant_fold_through=False, skip_signature_check=False)[源代码]¶
为函数(通常是来自 C 扩展的 C 函数)注册一个 polyfill 处理程序,以便在图中内联原始函数时使用该处理程序代替原始函数。
注意
仅在内联原始函数时才使用 polyfill 处理程序。当直接调用原始函数时,不会使用它。在急切模式下,装饰后的函数调用高性能 C 函数而不是 polyfill 处理程序。
polyfill 处理程序是一个函数,在内联原始函数时,它将被调用以代替原始函数。polyfill 处理程序应具有与原始函数相同的签名和行为。
- 参数
- 返回值
一个装饰器,为原始函数注册 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