jax.experimental.custom_dce.custom_dce.def_dce

custom_dce.def_dce(dce_rule)

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

参数:

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

返回类型:

Callable[[…], Any]