jax.experimental.custom_dce.custom_dce.def_dce

custom_dce.def_dce(dce_rule)

为此函数定义自定义 DCE 规则。

参数:

dce_rule (Callable[[...], Any]) – 一个函数，它接受 (a) 使用 static_argnums 指示为静态的任何参数，(b) 一个 bool 值的 Pytree (used_outs) 指示应该计算哪些输出，以及 (c) 原始函数的其余 (非静态) 参数。该规则应该返回一个与原始函数输出具有相同结构的 Pytree，但任何未使用的输出 (如 used_outs 所示) 都可以替换为 None。

返回类型:

Callable[[…], Any]