torch.compiler.substitute_in_graph

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

註冊函數的 polyfill 處理程序，通常是來自 C 擴展的 C 函數，用於在圖表中內聯原始函數時取代原始函數。

注意

polyfill 處理程序僅在內聯原始函數時使用。當直接調用原始函數時，它不會被使用。在 eager 模式下，裝飾後的函數調用效能高的 C 函數，而不是 polyfill 處理程序。

polyfill 處理器是一個函式，當內聯（inlining）原始函式時，將會呼叫此函式來取代原始函式。polyfill 處理器應該具有與原始函式相同的簽名 (signature) 和行為。

參數

• original_fn (callable) – 要註冊 polyfill 處理器的原始函式，通常是 C 函式。

• can_constant_fold_through (bool, optional) – 是否可以將 polyfill 處理器進行常數摺疊（constant fold）。也就是說，如果 polyfill 處理器是一個純函數 (pure function) 且其參數為常數，則在編譯期間可以對 polyfill 處理器的結果進行常數摺疊。預設值為 False。

• skip_signature_check (bool, optional) – 是否跳過原始函式和 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