functorch.compile.default_partition

functorch.compile.default_partition(joint_module, _joint_inputs, *, num_fwd_outputs)

此函式會分割 joint_module，其分割方式與在可呼叫物件的原始 .forward() 和 .backward() 中觀察到的行為非常相似，也就是說，產生的前向圖包含在傳遞給 aot_function() 的原始 .forward() 可呼叫物件中執行的那些運算子。

預設分割器會收集前向輸入和前向輸出之間的運算子。這有助於找到必須為反向傳遞儲存的張量。這些儲存的張量成為生成的前向圖的輸出。其餘的運算子則被放置在反向圖中。

警告

此 API 仍處於實驗階段，可能會有所變更。

參數

joint_module (fx.GraphModule) – 聯合前向和反向圖。這是 AOT Autograd 追蹤的結果。

回傳值

回傳生成的前向和反向 Fx 圖模組。