custom_vjp
和 nondiff_argnums
更新指南#
mattjj@ 2020 年 10 月 14 日
本文档假定您熟悉 jax.custom_vjp
,如用于 JAX 可变换 Python 函数的自定义导数规则笔记本中所述。
需要更新的内容#
在 JAX PR #4008 之后,传递给 custom_vjp
函数的 nondiff_argnums
的参数不能是 Tracer
(或 Tracer
的容器),这基本上意味着为了允许任意可转换的代码,nondiff_argnums
不应使用于数组值参数。相反,nondiff_argnums
应该仅用于非数组值,例如 Python 可调用对象或形状元组或字符串。
无论我们过去在哪里使用 nondiff_argnums
用于数组值,我们应该只将它们作为常规参数传递。在 bwd
规则中,我们需要为它们生成值,但我们可以只生成 None
值,以指示没有相应的梯度值。
例如,这是编写 clip_gradient
的旧方法,当 hi
和/或 lo
是来自某些 JAX 变换的 Tracer
时,这种方法将不起作用。
from functools import partial
import jax
@partial(jax.custom_vjp, nondiff_argnums=(0, 1))
def clip_gradient(lo, hi, x):
return x # identity function
def clip_gradient_fwd(lo, hi, x):
return x, None # no residual values to save
def clip_gradient_bwd(lo, hi, _, g):
return (jnp.clip(g, lo, hi),)
clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)
这是新的、很棒的方法,它支持任意变换
import jax
@jax.custom_vjp # no nondiff_argnums!
def clip_gradient(lo, hi, x):
return x # identity function
def clip_gradient_fwd(lo, hi, x):
return x, (lo, hi) # save lo and hi values as residuals
def clip_gradient_bwd(res, g):
lo, hi = res
return (None, None, jnp.clip(g, lo, hi)) # return None for lo and hi
clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)
如果您使用旧方法而不是新方法,在任何可能出错的情况下(即当 Tracer
传递到 nondiff_argnums
参数时),您都会收到响亮的错误提示。
这是一个我们实际上需要将 nondiff_argnums
与 custom_vjp
一起使用的情况
from functools import partial
import jax
@partial(jax.custom_vjp, nondiff_argnums=(0,))
def skip_app(f, x):
return f(x)
def skip_app_fwd(f, x):
return skip_app(f, x), None
def skip_app_bwd(f, _, g):
return (g,)
skip_app.defvjp(skip_app_fwd, skip_app_bwd)
解释#
将 Tracer
传递到 nondiff_argnums
参数总是存在 bug。虽然在某些情况下可以正确工作,但其他情况会导致复杂且令人困惑的错误消息。
bug 的本质是 nondiff_argnums
的实现方式非常像词法闭包。但当时 Tracers 上的词法闭包并不打算与 custom_jvp
/custom_vjp
一起使用。以这种方式实现 nondiff_argnums
是一个错误!
PR #4008 修复了 custom_jvp
和 custom_vjp
的所有词法闭包问题。 哇呼!也就是说,现在 custom_jvp
和 custom_vjp
函数和规则可以随意闭包 Tracers。对于所有非自动微分变换,一切都将正常工作。对于自动微分变换,我们将收到一个明确的错误消息,说明为什么我们不能对 custom_jvp
或 custom_vjp
闭包的值进行微分
检测到对 custom_jvp 函数相对于闭包值的微分。这是不支持的,因为自定义 JVP 规则仅指定如何相对于显式输入参数来区分 custom_jvp 函数。
尝试将闭包值作为参数传递到 custom_jvp 函数中,并调整 custom_jvp 规则。
以这种方式加强和增强 custom_jvp
和 custom_vjp
时,我们发现允许 custom_vjp
在其 nondiff_argnums
中接受 Tracers 将需要大量的簿记工作:我们需要重写用户的 fwd
函数以将值作为残差返回,并重写用户的 bwd
函数以将它们作为正常残差接受(而不是像 nondiff_argnums
那样将它们作为特殊的引导参数接受)。这看起来也许是可管理的,直到您考虑我们如何处理任意 pytrees!而且,这种复杂性是不必要的:如果用户代码像处理常规参数和残差一样处理类似数组的非可微参数,那么一切都已经正常工作了。(在 #4039 之前,JAX 可能会抱怨在自动微分中涉及整数值输入和输出,但在 #4039 之后,这些将正常工作!)
与 custom_vjp
不同,很容易使 custom_jvp
与作为 Tracers 的 nondiff_argnums
参数一起工作。因此,这些更新只需要在 custom_vjp
中进行。