functorch.compile.default_partition

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

以与原始.forward()和.backward()中观察到的行为非常相似的方式对joint_module进行分区，即，生成的前向图包含在传递给aot_function()的原始.forward()可调用对象中执行的那些操作符。

默认的分区器收集位于前向输入和前向输出之间的操作符。这有助于找到需要为反向传播暂存的张量。这些暂存的张量成为生成的前向图的输出。然后，剩余的操作符被放置在反向图中。

警告

此API是实验性的，可能会发生变化。

Parameters

joint_module (fx.GraphModule) – 联合的前向和后向图。这是AOT Autograd追踪的结果。

Returns

返回生成的前向和后向Fx图模块。