torch.compiler.allow_in_graph

torch.compiler.allow_in_graph(fn)

指示編譯器前端（Dynamo）在遇到函式時跳過其符號內省，而是直接將其寫入圖。.

如果您正在使用 torch.compile()（使用 backend=”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() 的單個圖。