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+Inductor)
無裝飾器	追蹤內部	追蹤內部
allow_in_graph	不透明可呼叫物件	追蹤內部
自定義運算元	不透明可呼叫物件	不透明可呼叫物件

allow_in_graph() 的一個常見用例是作為編譯器前端的逃生艙：如果您知道該函式對於編譯堆疊的下游元件（AOTAutograd 和 Inductor）來說是有效的，但存在 Dynamo 的錯誤阻止它正確地符號內省該函式（或者如果您的程式碼是用 C/C++ 編寫，因此無法使用 Dynamo 進行內省），那麼可以使用 allow_in_graph() 裝飾該函式以繞過 Dynamo。

我們要求 fn 遵循以下限制。不遵守將導致未定義行為。

• fn 的輸入必須是 FX 圖中可代理 (Proxy-able) 的型別。有效型別包括：Tensor/int/bool/float/None/List[Tensor?]/List[int?]/List[float?] Tuple[Tensor?, …]/Tuple[int?, …]/Tuple[float?, …]/torch.dtype/torch.device

• fn 的輸出必須是 FX 圖中可代理 (Proxy-able) 的型別（參見上一條）。

• fn 內部使用的所有 Tensor 必須直接作為輸入傳遞給 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() 的單一圖。