jax.custom_gradient

jax.custom_gradient(fun)

用于定义自定义 VJP 规则（即自定义梯度）的便捷函数。

虽然定义自定义 VJP 规则的规范方法是通过 jax.custom_vjp，但 custom_gradient 便捷封装遵循 TensorFlow 的 tf.custom_gradient API。这里的区别在于 custom_gradient 可以用作一个函数的装饰器，该函数返回原始值（表示要微分的数学函数的输出）和 VJP（梯度）函数。请参阅 https://tensorflowcn.cn/api_docs/python/tf/custom_gradient。

如果待微分的数学函数具有 Haskell 风格的签名 a -> b，那么 Python 可调用对象 fun 应具有签名 a -> (b, CT b --o CT a)，其中我们使用 CT x 来表示 x 的余切类型，--o 箭头表示线性函数。请参见下方的示例。也就是说，fun 应返回一个元组，其中第一个元素表示待微分数学函数的值，第二个元素是将在反向模式自动微分的反向传播中调用的函数（即“自定义梯度”函数）。

fun 输出的第二个元素所返回的函数可以捕获在评估待微分函数时计算的中间值。也就是说，可以使用词法闭包来在反向模式自动微分的前向传播和反向传播之间共享工作。但是，它不能执行依赖于捕获的中间值或其余切参数值的 Python 控制流；如果函数包含此类控制流，则会引发错误。

参数:

fun – 一个 Python 可调用对象，用于指定待微分的数学函数及其反向模式微分规则。它应返回一个元组，由一个输出值和一个表示自定义梯度函数的 Python 可调用对象组成。

返回:

一个 Python 可调用对象，它接受与 fun 相同的参数，并返回 fun 输出元组的第一个元素所指定的输出值。

例如

>>> @jax.custom_gradient
... def f(x):
...   return x ** 2, lambda g: (g * x,)
...
>>> print(f(3.))
9.0
>>> print(jax.grad(f)(3.))
3.0

一个带有两个参数的函数的示例，因此 VJP 函数必须返回一个长度为二的元组

>>> @jax.custom_gradient
... def f(x, y):
...   return x * y, lambda g: (g * y, g * x)
...
>>> print(f(3., 4.))
12.0
>>> print(jax.grad(f, argnums=(0, 1))(3., 4.))
(Array(4., dtype=float32, weak_type=True), Array(3., dtype=float32, weak_type=True))