functorch.compile.default_partition

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

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

默认分区器收集正向输入和正向输出之间的运算符。这有助于找到必须为反向传递存储的张量。这些存储的张量成为生成的正向图的输出。然后，其余运算符将放置在反向图中。

警告

此 API 处于实验阶段，可能会发生变化。

参数

joint_module (fx.GraphModule) – 联合的正向和反向图。这是 AOT 自动微分跟踪的结果。

返回值

返回生成的正向和反向 Fx 图模块。