jax.自定义梯度

jax.custom_gradient(fun)

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

虽然定义自定义 VJP 规则的规范方式是通过 jax.custom_vjp，但 custom_gradient 便捷包装器遵循 TensorFlow 的 tf.custom_gradient API。这里的区别在于 custom_gradient 可以用作装饰器，作用于一个返回原始值（表示要微分的数学函数的输出）和 VJP（梯度）函数的函数。参见 https://www.tensorflow.org/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可调用对象组成的对。

返回:

一个接受与 fun 相同参数的 Python 可调用对象，并返回由 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))