自定义导数规则

在 JAX 中有两种方法可以定义微分规则

1. 使用 jax.custom_jvp 和 jax.custom_vjp 为已经可以 JAX 变换的 Python 函数定义自定义微分规则；以及

2. 定义新的 core.Primitive 实例及其所有转换规则，例如调用其他系统（如求解器、模拟器或通用数值计算系统）中的函数。

本笔记本主要介绍 #1。要阅读有关 #2 的内容，请参阅 添加原语的笔记本。

有关 JAX 自动微分 API 的介绍，请参阅 自动微分食谱。本笔记本假设您已熟悉 jax.jvp 和 jax.grad，以及 JVP 和 VJP 的数学含义。

摘要

使用 jax.custom_jvp 自定义 JVP

import jax.numpy as jnp
from jax import custom_jvp

@custom_jvp
def f(x, y):
  return jnp.sin(x) * y

@f.defjvp
def f_jvp(primals, tangents):
  x, y = primals
  x_dot, y_dot = tangents
  primal_out = f(x, y)
  tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
  return primal_out, tangent_out

from jax import jvp, grad

print(f(2., 3.))
y, y_dot = jvp(f, (2., 3.), (1., 0.))
print(y)
print(y_dot)
print(grad(f)(2., 3.))

2.7278922
2.7278922
-1.2484405
-1.2484405

# Equivalent alternative using the defjvps convenience wrapper

@custom_jvp
def f(x, y):
  return jnp.sin(x) * y

f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,
          lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)

print(f(2., 3.))
y, y_dot = jvp(f, (2., 3.), (1., 0.))
print(y)
print(y_dot)
print(grad(f)(2., 3.))

2.7278922
2.7278922
-1.2484405
-1.2484405

使用 jax.custom_vjp 自定义 VJP

from jax import custom_vjp

@custom_vjp
def f(x, y):
  return jnp.sin(x) * y

def f_fwd(x, y):
# Returns primal output and residuals to be used in backward pass by f_bwd.
  return f(x, y), (jnp.cos(x), jnp.sin(x), y)

def f_bwd(res, g):
  cos_x, sin_x, y = res # Gets residuals computed in f_fwd
  return (cos_x * g * y, sin_x * g)

f.defvjp(f_fwd, f_bwd)

print(grad(f)(2., 3.))

-1.2484405

示例问题

为了了解 jax.custom_jvp 和 jax.custom_vjp 用于解决哪些问题，让我们来看几个例子。下一节将对 jax.custom_jvp 和 jax.custom_vjp API 进行更深入的介绍。

数值稳定性

jax.custom_jvp 的一个应用是提高微分的数值稳定性。

假设我们想编写一个名为 log1pexp 的函数，它计算 \(x \mapsto \log ( 1 + e^x )\)。我们可以使用 jax.numpy 来编写它。

def log1pexp(x):
  return jnp.log(1. + jnp.exp(x))

log1pexp(3.)

Array(3.0485873, dtype=float32, weak_type=True)

由于它是用 jax.numpy 编写的，因此它是可 JAX 转换的。

from jax import jit, grad, vmap

print(jit(log1pexp)(3.))
print(jit(grad(log1pexp))(3.))
print(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))

3.0485873
0.95257413
[0.5       0.7310586 0.8807971]

但是这里潜藏着一个数值稳定性问题。

print(grad(log1pexp)(100.))

nan

这似乎不对！毕竟，\(x \mapsto \log (1 + e^x)\) 的导数是 \(x \mapsto \frac{e^x}{1 + e^x}\)，因此对于 \(x\) 的较大值，我们期望该值为约 1。

我们可以通过查看梯度计算的 jaxpr 来更深入地了解正在发生的事情。

from jax import make_jaxpr

make_jaxpr(grad(log1pexp))(100.)

{ lambda ; a:f32[]. let
    b:f32[] = exp a
    c:f32[] = add 1.0 b
    _:f32[] = log c
    d:f32[] = div 1.0 c
    e:f32[] = mul d b
  in (e,) }

逐步执行 jaxpr 的评估方式，我们可以看到最后一行将涉及乘以浮点数学将四舍五入为 0 和 \(\infty\) 的值，这绝不是一个好主意。也就是说，我们实际上正在为较大的 x 评估 lambda x: (1 / (1 + jnp.exp(x))) * jnp.exp(x)，这实际上变成了 0. * jnp.inf。

与其生成如此大和小的值，并希望进行浮点数无法始终提供的抵消，不如直接将导数函数表示为更具数值稳定性的程序。特别是，我们可以编写一个程序，该程序更接近于评估相等的数学表达式 \(1 - \frac{1}{1 + e^x}\)，而无需进行抵消。

这个问题很有趣，因为即使我们对 log1pexp 的定义已经可以进行 JAX 微分（并使用 jit、vmap 等进行转换），但我们对将标准自动微分规则应用于构成 log1pexp 的原语并将结果组合起来的结果并不满意。相反，我们希望指定整个函数 log1pexp 如何作为整体进行微分，从而更好地安排这些指数。

这是已经可 JAX 转换的 Python 函数的自定义导数规则的一个应用：指定复合函数如何进行微分，同时仍然使用其原始 Python 定义进行其他转换（如 jit、vmap 等）。

这是一个使用 jax.custom_jvp 的解决方案。

from jax import custom_jvp

@custom_jvp
def log1pexp(x):
  return jnp.log(1. + jnp.exp(x))

@log1pexp.defjvp
def log1pexp_jvp(primals, tangents):
  x, = primals
  x_dot, = tangents
  ans = log1pexp(x)
  ans_dot = (1 - 1/(1 + jnp.exp(x))) * x_dot
  return ans, ans_dot

print(grad(log1pexp)(100.))

1.0

print(jit(log1pexp)(3.))
print(jit(grad(log1pexp))(3.))
print(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))

3.0485873
0.95257413
[0.5       0.7310586 0.8807971]

这是一个 defjvps 辅助包装器，用于表达相同的内容。

@custom_jvp
def log1pexp(x):
  return jnp.log(1. + jnp.exp(x))

log1pexp.defjvps(lambda t, ans, x: (1 - 1/(1 + jnp.exp(x))) * t)

print(grad(log1pexp)(100.))
print(jit(log1pexp)(3.))
print(jit(grad(log1pexp))(3.))
print(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))

1.0
3.0485873
0.95257413
[0.5       0.7310586 0.8807971]

强制执行微分约定

相关的应用是强制执行微分约定，可能是在边界处。

考虑函数 \(f : \mathbb{R}_+ \to \mathbb{R}_+\)，其中 \(f(x) = \frac{x}{1 + \sqrt{x}}\)，我们取 \(\mathbb{R}_+ = [0, \infty)\)。我们可以将 \(f\) 实现为如下程序。

def f(x):
  return x / (1 + jnp.sqrt(x))

作为 \(\mathbb{R}\)（完整实数线）上的数学函数，\(f\) 在零处不可微（因为定义导数的极限从左侧不存在）。相应地，自动微分会产生 nan 值。

print(grad(f)(0.))

nan

但在数学上，如果我们将 \(f\) 视为 \(\mathbb{R}_+\) 上的函数，那么它在 0 处可微[Rudin 的《数学分析原理》定义 5.1，或 Tao 的《分析 I》第 3 版定义 10.1.1 和示例 10.1.6]。或者，我们可以说按照约定，我们希望考虑右侧的方向导数。因此，Python 函数 grad(f) 在 0.0 处返回一个合理的数值，即 1.0。默认情况下，JAX 的微分机制假设所有函数都定义在 \(\mathbb{R}\) 上，因此这里不会产生 1.0。

我们可以使用自定义 JVP 规则！特别是，我们可以根据 \(\mathbb{R}_+\) 上的导数函数 \(x \mapsto \frac{\sqrt{x} + 2}{2(\sqrt{x} + 1)^2}\) 定义 JVP 规则。

@custom_jvp
def f(x):
  return x / (1 + jnp.sqrt(x))

@f.defjvp
def f_jvp(primals, tangents):
  x, = primals
  x_dot, = tangents
  ans = f(x)
  ans_dot = ((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * x_dot
  return ans, ans_dot

print(grad(f)(0.))

1.0

这是辅助包装器版本。

@custom_jvp
def f(x):
  return x / (1 + jnp.sqrt(x))

f.defjvps(lambda t, ans, x: ((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * t)

print(grad(f)(0.))

1.0

梯度裁剪

在某些情况下，我们希望表达数学微分计算，而在其他情况下，我们甚至可能希望从数学中脱离一步以调整自动微分执行的计算。一个典型的例子是反向模式梯度裁剪。

对于梯度裁剪，我们可以使用 jnp.clip 以及 jax.custom_vjp 反向模式专用规则。

from functools import partial
from jax import custom_vjp

@custom_vjp
def clip_gradient(lo, hi, x):
  return x  # identity function

def clip_gradient_fwd(lo, hi, x):
  return x, (lo, hi)  # save bounds as residuals

def clip_gradient_bwd(res, g):
  lo, hi = res
  return (None, None, jnp.clip(g, lo, hi))  # use None to indicate zero cotangents for lo and hi

clip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)

import matplotlib.pyplot as plt
from jax import vmap

t = jnp.linspace(0, 10, 1000)

plt.plot(jnp.sin(t))
plt.plot(vmap(grad(jnp.sin))(t))

[<matplotlib.lines.Line2D at 0x7feef0acd450>]

def clip_sin(x):
  x = clip_gradient(-0.75, 0.75, x)
  return jnp.sin(x)

plt.plot(clip_sin(t))
plt.plot(vmap(grad(clip_sin))(t))

[<matplotlib.lines.Line2D at 0x7feef09ca980>]

Python 调试

另一个受开发工作流而非数值驱动的应用是在反向模式自动微分的反向传播过程中设置 pdb 调试器跟踪。

在尝试跟踪 nan 运行时错误的根源或仔细检查正在传播的余切（梯度）值时，在对应于原始计算中特定点的反向传播中插入调试器可能很有用。您可以使用 jax.custom_vjp 来实现。

我们将把示例推迟到下一节。

迭代实现的隐函数微分

这个例子在数学方面深入探讨！

jax.custom_vjp 的另一个应用是反向模式微分，用于可 JAX 转换（通过 jit、vmap 等）但由于某些原因无法有效 JAX 微分的函数，可能是因为它们涉及 lax.while_loop。（无法生成有效计算 XLA HLO While 循环的反向模式导数的 XLA HLO 程序，因为这需要一个具有无限内存使用的程序，这在 XLA HLO 中无法表达，至少在没有通过进/出馈进行副作用交互的情况下是不可能的。）

例如，考虑这个 fixed_point 例程，它通过在 while_loop 中迭代应用函数来计算不动点。

from jax.lax import while_loop

def fixed_point(f, a, x_guess):
  def cond_fun(carry):
    x_prev, x = carry
    return jnp.abs(x_prev - x) > 1e-6

  def body_fun(carry):
    _, x = carry
    return x, f(a, x)

  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))
  return x_star

这是一种数值求解方程 \(x = f(a, x)\) 的迭代过程，求解 \(x\)，通过迭代 \(x_{t+1} = f(a, x_t)\) 直到 \(x_{t+1}\) 足够接近 \(x_t\)。结果 \(x^*\) 取决于参数 \(a\)，因此我们可以认为存在一个由方程 \(x = f(a, x)\) 隐式定义的函数 \(a \mapsto x^*(a)\)。

我们可以使用 fixed_point 来运行迭代过程直至收敛，例如运行牛顿法来计算平方根，同时仅执行加法、乘法和除法。

def newton_sqrt(a):
  update = lambda a, x: 0.5 * (x + a / x)
  return fixed_point(update, a, a)

print(newton_sqrt(2.))

1.4142135

我们也可以对函数使用 vmap 或 jit。

print(jit(vmap(newton_sqrt))(jnp.array([1., 2., 3., 4.])))

[1.        1.4142135 1.7320509 2.       ]

我们不能应用反向模式自动微分，因为存在 while_loop，但事实证明我们也不希望这样做：与其遍历 fixed_point 及其所有迭代的实现进行微分，不如利用数学结构来做一些更节省内存（在本例中也更节省浮点运算）的事情！相反，我们可以使用隐函数定理[Bertsekas 的《非线性规划》第 2 版的命题 A.25]，该定理保证（在某些条件下）我们即将使用的数学对象的存在。从本质上讲，我们在解处线性化并迭代求解这些线性方程，以计算我们想要的导数。

再次考虑方程 \(x = f(a, x)\) 和函数 \(x^*\)。我们希望评估向量-雅可比乘积，如 \(v^\mathsf{T} \mapsto v^\mathsf{T} \partial x^*(a_0)\)。

至少在我们想要求导的点\(a_0\)周围的一个开邻域内，假设方程\(x^*(a) = f(a, x^*(a))\)对所有\(a\)都成立。由于等式两边作为\(a\)的函数是相等的，它们的导数也必须相等，所以让我们对等式两边求导

\(\qquad \partial x^*(a) = \partial_0 f(a, x^*(a)) + \partial_1 f(a, x^*(a)) \partial x^*(a)\).

令\(A = \partial_1 f(a_0, x^*(a_0))\)和\(B = \partial_0 f(a_0, x^*(a_0))\)，我们可以将我们想要计算的量更简单地写成

\(\qquad \partial x^*(a_0) = B + A \partial x^*(a_0)\),

或者，通过重新排列，

\(\qquad \partial x^*(a_0) = (I - A)^{-1} B\).

这意味着我们可以计算像这样的向量-雅可比乘积

\(\qquad v^\mathsf{T} \partial x^*(a_0) = v^\mathsf{T} (I - A)^{-1} B = w^\mathsf{T} B\),

其中\(w^\mathsf{T} = v^\mathsf{T} (I - A)^{-1}\)，或者等价地\(w^\mathsf{T} = v^\mathsf{T} + w^\mathsf{T} A\)，或者等价地\(w^\mathsf{T}\)是映射\(u^\mathsf{T} \mapsto v^\mathsf{T} + u^\mathsf{T} A\)的不动点。最后一个表述为我们提供了一种方法，可以用对fixed_point的调用来编写fixed_point的VJP！此外，在将\(A\)和\(B\)展开后，我们可以看到我们只需要在\((a_0, x^*(a_0))\)处评估\(f\)的VJP。

以下是要点

from jax import vjp

@partial(custom_vjp, nondiff_argnums=(0,))
def fixed_point(f, a, x_guess):
  def cond_fun(carry):
    x_prev, x = carry
    return jnp.abs(x_prev - x) > 1e-6

  def body_fun(carry):
    _, x = carry
    return x, f(a, x)

  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))
  return x_star

def fixed_point_fwd(f, a, x_init):
  x_star = fixed_point(f, a, x_init)
  return x_star, (a, x_star)

def fixed_point_rev(f, res, x_star_bar):
  a, x_star = res
  _, vjp_a = vjp(lambda a: f(a, x_star), a)
  a_bar, = vjp_a(fixed_point(partial(rev_iter, f),
                             (a, x_star, x_star_bar),
                             x_star_bar))
  return a_bar, jnp.zeros_like(x_star)

def rev_iter(f, packed, u):
  a, x_star, x_star_bar = packed
  _, vjp_x = vjp(lambda x: f(a, x), x_star)
  return x_star_bar + vjp_x(u)[0]

fixed_point.defvjp(fixed_point_fwd, fixed_point_rev)

print(newton_sqrt(2.))

1.4142135

print(grad(newton_sqrt)(2.))
print(grad(grad(newton_sqrt))(2.))

0.35355338
-0.088388346

我们可以通过对jnp.sqrt求导来检查我们的答案，它使用了完全不同的实现

print(grad(jnp.sqrt)(2.))
print(grad(grad(jnp.sqrt))(2.))

0.35355338
-0.08838835

这种方法的一个限制是参数f不能封闭任何参与微分的数值。也就是说，您可能会注意到我们将参数a显式地放在了fixed_point的参数列表中。对于这种情况，请考虑使用低级原语lax.custom_root，它允许使用自定义寻根函数对封闭变量进行求导。

jax.custom_jvp和jax.custom_vjp API的基本用法

使用jax.custom_jvp定义前向模式（以及间接的反向模式）规则

这是一个使用jax.custom_jvp的典型基本示例，其中注释使用了类似 Haskell 的类型签名

from jax import custom_jvp
import jax.numpy as jnp

# f :: a -> b
@custom_jvp
def f(x):
  return jnp.sin(x)

# f_jvp :: (a, T a) -> (b, T b)
def f_jvp(primals, tangents):
  x, = primals
  t, = tangents
  return f(x), jnp.cos(x) * t

f.defjvp(f_jvp)

<function __main__.f_jvp(primals, tangents)>

from jax import jvp

print(f(3.))

y, y_dot = jvp(f, (3.,), (1.,))
print(y)
print(y_dot)

0.14112
0.14112
-0.9899925

换句话说，我们从一个原始函数f开始，它接受类型为a的输入并产生类型为b的输出。我们将其与一个 JVP 规则函数f_jvp关联起来，该函数接受一对输入，表示类型为a的原始输入和类型为T a的相应切线输入，并产生一对输出，表示类型为b的原始输出和类型为T b的切线输出。切线输出应该是切线输入的线性函数。

您还可以使用f.defjvp作为装饰器，例如

@custom_jvp
def f(x):
  ...

@f.defjvp
def f_jvp(primals, tangents):
  ...

即使我们只定义了 JVP 规则而没有定义 VJP 规则，我们也可以对f使用前向和反向模式微分。JAX 将自动转置我们自定义 JVP 规则中对切线值的线性计算，计算 VJP 的效率与我们手动编写规则一样高。

from jax import grad

print(grad(f)(3.))
print(grad(grad(f))(3.))

-0.9899925
-0.14112

为了使自动转置工作，JVP 规则的输出切线必须作为输入切线的函数是线性的。否则将引发转置错误。

多个参数的工作方式如下

@custom_jvp
def f(x, y):
  return x ** 2 * y

@f.defjvp
def f_jvp(primals, tangents):
  x, y = primals
  x_dot, y_dot = tangents
  primal_out = f(x, y)
  tangent_out = 2 * x * y * x_dot + x ** 2 * y_dot
  return primal_out, tangent_out

print(grad(f)(2., 3.))

12.0

defjvps便捷包装器允许我们分别为每个参数定义一个 JVP，并且结果将分别计算然后求和。

@custom_jvp
def f(x):
  return jnp.sin(x)

f.defjvps(lambda t, ans, x: jnp.cos(x) * t)

print(grad(f)(3.))

-0.9899925

这是一个带有多个参数的defjvps示例

@custom_jvp
def f(x, y):
  return x ** 2 * y

f.defjvps(lambda x_dot, primal_out, x, y: 2 * x * y * x_dot,
          lambda y_dot, primal_out, x, y: x ** 2 * y_dot)

print(grad(f)(2., 3.))
print(grad(f, 0)(2., 3.))  # same as above
print(grad(f, 1)(2., 3.))

12.0
12.0
4.0

作为简写，使用defjvps，您可以传递一个None值来指示特定参数的 JVP 为零。

@custom_jvp
def f(x, y):
  return x ** 2 * y

f.defjvps(lambda x_dot, primal_out, x, y: 2 * x * y * x_dot,
          None)

print(grad(f)(2., 3.))
print(grad(f, 0)(2., 3.))  # same as above
print(grad(f, 1)(2., 3.))

12.0
12.0
0.0

使用关键字参数调用jax.custom_jvp函数，或编写具有默认参数的jax.custom_jvp函数定义都是允许的，只要它们可以根据标准库inspect.signature机制检索的函数签名明确地映射到位置参数即可。

当您不执行微分时，函数f就像没有被jax.custom_jvp装饰一样被调用。

@custom_jvp
def f(x):
  print('called f!')  # a harmless side-effect
  return jnp.sin(x)

@f.defjvp
def f_jvp(primals, tangents):
  print('called f_jvp!')  # a harmless side-effect
  x, = primals
  t, = tangents
  return f(x), jnp.cos(x) * t

from jax import vmap, jit

print(f(3.))

called f!
0.14112

print(vmap(f)(jnp.arange(3.)))
print(jit(f)(3.))

called f!
[0.         0.84147096 0.9092974 ]
called f!
0.14112

自定义 JVP 规则在微分过程中被调用，无论是前向还是反向。

y, y_dot = jvp(f, (3.,), (1.,))
print(y_dot)

called f_jvp!
called f!
-0.9899925

print(grad(f)(3.))

called f_jvp!
called f!
-0.9899925

请注意，f_jvp调用f来计算原始输出。在高阶微分的上下文中，每次应用微分变换都将使用自定义 JVP 规则，当且仅当该规则调用原始f来计算原始输出时。（这代表了一种基本权衡，我们不能在规则中利用f评估的中间值，并且还让规则适用于所有高阶微分。）

grad(grad(f))(3.)

called f_jvp!
called f_jvp!
called f!

Array(-0.14112, dtype=float32, weak_type=True)

您可以在jax.custom_jvp中使用 Python 控制流。

@custom_jvp
def f(x):
  if x > 0:
    return jnp.sin(x)
  else:
    return jnp.cos(x)

@f.defjvp
def f_jvp(primals, tangents):
  x, = primals
  x_dot, = tangents
  ans = f(x)
  if x > 0:
    return ans, 2 * x_dot
  else:
    return ans, 3 * x_dot

print(grad(f)(1.))
print(grad(f)(-1.))

2.0
3.0

使用jax.custom_vjp定义自定义的反向模式规则

虽然jax.custom_jvp足以控制前向模式和（通过 JAX 的自动转置）反向模式微分行为，但在某些情况下，我们可能希望直接控制 VJP 规则，例如在上文介绍的后两个示例问题中。我们可以使用jax.custom_vjp来做到这一点。

from jax import custom_vjp
import jax.numpy as jnp

# f :: a -> b
@custom_vjp
def f(x):
  return jnp.sin(x)

# f_fwd :: a -> (b, c)
def f_fwd(x):
  return f(x), jnp.cos(x)

# f_bwd :: (c, CT b) -> CT a
def f_bwd(cos_x, y_bar):
  return (cos_x * y_bar,)

f.defvjp(f_fwd, f_bwd)

from jax import grad

print(f(3.))
print(grad(f)(3.))

0.14112
-0.9899925

换句话说，我们再次从一个原始函数f开始，它接受类型为a的输入并产生类型为b的输出。我们将其与两个函数f_fwd和f_bwd关联起来，它们分别描述了如何执行反向模式自动微分的前向和后向传递。

函数f_fwd描述了前向传递，不仅包括原始计算，还包括哪些值需要保存以供后向传递使用。它的输入签名与原始函数f的输入签名相同，因为它接受类型为a的原始输入。但作为输出，它产生一对，其中第一个元素是原始输出b，第二个元素是类型为c的任何“残差”数据，用于存储以供后向传递使用。（此第二个输出类似于PyTorch 的 save_for_backward 机制。）

函数f_bwd描述了后向传递。它接受两个输入，第一个是f_fwd生成的类型为c的残差数据，第二个是对应于原始函数输出的类型为CT b的输出余切。它产生类型为CT a的输出，表示对应于原始函数输入的余切。特别是，f_bwd的输出必须是长度等于原始函数参数数量的序列（例如元组）。

所以多个参数的工作方式如下

from jax import custom_vjp

@custom_vjp
def f(x, y):
  return jnp.sin(x) * y

def f_fwd(x, y):
  return f(x, y), (jnp.cos(x), jnp.sin(x), y)

def f_bwd(res, g):
  cos_x, sin_x, y = res
  return (cos_x * g * y, sin_x * g)

f.defvjp(f_fwd, f_bwd)

print(grad(f)(2., 3.))

-1.2484405

使用关键字参数调用jax.custom_vjp函数，或编写具有默认参数的jax.custom_vjp函数定义都是允许的，只要它们可以根据标准库inspect.signature机制检索的函数签名明确地映射到位置参数即可。

与jax.custom_jvp一样，如果未应用微分，则不会调用由f_fwd和f_bwd组成的自定义 VJP 规则。如果函数被求值或使用jit、vmap或其他非微分变换进行变换，则只调用f。

@custom_vjp
def f(x):
  print("called f!")
  return jnp.sin(x)

def f_fwd(x):
  print("called f_fwd!")
  return f(x), jnp.cos(x)

def f_bwd(cos_x, y_bar):
  print("called f_bwd!")
  return (cos_x * y_bar,)

f.defvjp(f_fwd, f_bwd)

print(f(3.))

called f!
0.14112

print(grad(f)(3.))

called f_fwd!
called f!
called f_bwd!
-0.9899925

y, f_vjp = vjp(f, 3.)
print(y)

called f_fwd!
called f!
0.14112

print(f_vjp(1.))

called f_bwd!
(Array(-0.9899925, dtype=float32, weak_type=True),)

**无法对**jax.custom_vjp**函数使用前向模式自动微分**，这将引发错误。

from jax import jvp

try:
  jvp(f, (3.,), (1.,))
except TypeError as e:
  print('ERROR! {}'.format(e))

called f_fwd!
called f!
ERROR! can't apply forward-mode autodiff (jvp) to a custom_vjp function.

如果您想使用前向和反向模式，请改用jax.custom_jvp。

我们可以将jax.custom_vjp与pdb结合使用，以便在后向传递中插入调试器跟踪。

import pdb

@custom_vjp
def debug(x):
  return x  # acts like identity

def debug_fwd(x):
  return x, x

def debug_bwd(x, g):
  pdb.set_trace()
  return g

debug.defvjp(debug_fwd, debug_bwd)

def foo(x):
  y = x ** 2
  y = debug(y)  # insert pdb in corresponding backward pass step
  return jnp.sin(y)

jax.grad(foo)(3.)

> <ipython-input-113-b19a2dc1abf7>(12)debug_bwd()
-> return g
(Pdb) p x
Array(9., dtype=float32)
(Pdb) p g
Array(-0.91113025, dtype=float32)
(Pdb) q

更多功能和细节

使用list / tuple / dict容器（和其他pytrees）

您应该期望标准的 Python 容器（如列表、元组、命名元组和字典）以及它们的嵌套版本都能正常工作。一般来说，任何pytrees都是允许的，只要它们根据类型约束保持一致即可。

这是一个使用jax.custom_jvp的虚构示例

from collections import namedtuple
Point = namedtuple("Point", ["x", "y"])

@custom_jvp
def f(pt):
  x, y = pt.x, pt.y
  return {'a': x ** 2,
          'b': (jnp.sin(x), jnp.cos(y))}

@f.defjvp
def f_jvp(primals, tangents):
  pt, = primals
  pt_dot, =  tangents
  ans = f(pt)
  ans_dot = {'a': 2 * pt.x * pt_dot.x,
             'b': (jnp.cos(pt.x) * pt_dot.x, -jnp.sin(pt.y) * pt_dot.y)}
  return ans, ans_dot

def fun(pt):
  dct = f(pt)
  return dct['a'] + dct['b'][0]

pt = Point(1., 2.)

print(f(pt))

{'a': 1.0, 'b': (Array(0.84147096, dtype=float32, weak_type=True), Array(-0.41614684, dtype=float32, weak_type=True))}

print(grad(fun)(pt))

Point(x=Array(2.5403023, dtype=float32, weak_type=True), y=Array(0., dtype=float32, weak_type=True))

以及一个使用jax.custom_vjp的类似构造示例

@custom_vjp
def f(pt):
  x, y = pt.x, pt.y
  return {'a': x ** 2,
          'b': (jnp.sin(x), jnp.cos(y))}

def f_fwd(pt):
  return f(pt), pt

def f_bwd(pt, g):
  a_bar, (b0_bar, b1_bar) = g['a'], g['b']
  x_bar = 2 * pt.x * a_bar + jnp.cos(pt.x) * b0_bar
  y_bar = -jnp.sin(pt.y) * b1_bar
  return (Point(x_bar, y_bar),)

f.defvjp(f_fwd, f_bwd)

def fun(pt):
  dct = f(pt)
  return dct['a'] + dct['b'][0]

pt = Point(1., 2.)

print(f(pt))

{'a': 1.0, 'b': (Array(0.84147096, dtype=float32, weak_type=True), Array(-0.41614684, dtype=float32, weak_type=True))}

print(grad(fun)(pt))

Point(x=Array(2.5403023, dtype=float32, weak_type=True), y=Array(-0., dtype=float32, weak_type=True))

处理不可微参数

某些用例，例如最后的示例问题，需要将不可微参数（例如函数值参数）传递给具有自定义微分规则的函数，并将这些参数也传递给规则本身。在fixed_point的情况下，函数参数f就是这样一个不可微参数。类似的情况也出现在jax.experimental.odeint中。

jax.custom_jvp 与 nondiff_argnums

使用可选的nondiff_argnums参数传递给jax.custom_jvp来指示这些参数。以下是一个使用jax.custom_jvp的示例

from functools import partial

@partial(custom_jvp, nondiff_argnums=(0,))
def app(f, x):
  return f(x)

@app.defjvp
def app_jvp(f, primals, tangents):
  x, = primals
  x_dot, = tangents
  return f(x), 2. * x_dot

print(app(lambda x: x ** 3, 3.))

27.0

print(grad(app, 1)(lambda x: x ** 3, 3.))

2.0

请注意这里的一个陷阱：无论这些参数出现在参数列表中的哪个位置，它们都将被放置在相应JVP规则签名中的开头。以下还有另一个示例

@partial(custom_jvp, nondiff_argnums=(0, 2))
def app2(f, x, g):
  return f(g((x)))

@app2.defjvp
def app2_jvp(f, g, primals, tangents):
  x, = primals
  x_dot, = tangents
  return f(g(x)), 3. * x_dot

print(app2(lambda x: x ** 3, 3., lambda y: 5 * y))

3375.0

print(grad(app2, 1)(lambda x: x ** 3, 3., lambda y: 5 * y))

3.0

jax.custom_vjp 与 nondiff_argnums

jax.custom_vjp也存在类似的选项，并且同样地，约定是将不可微参数作为_bwd规则的第一个参数传递，无论它们出现在原始函数签名中的哪个位置。_fwd规则的签名保持不变 - 它与原始函数的签名相同。以下是一个示例

@partial(custom_vjp, nondiff_argnums=(0,))
def app(f, x):
  return f(x)

def app_fwd(f, x):
  return f(x), x

def app_bwd(f, x, g):
  return (5 * g,)

app.defvjp(app_fwd, app_bwd)

print(app(lambda x: x ** 2, 4.))

16.0

print(grad(app, 1)(lambda x: x ** 2, 4.))

5.0

请参阅上面的fixed_point以获取另一个使用示例。

不需要对 nondiff_argnums与数组值参数一起使用，例如具有整数数据类型的值。相反，nondiff_argnums应该只用于不对应于JAX类型的参数值（本质上不对应于数组类型），例如Python可调用对象或字符串。如果JAX检测到由nondiff_argnums指示的参数包含JAX Tracer，则会引发错误。上面的clip_gradient函数就是一个不将nondiff_argnums用于整数数据类型数组参数的良好示例。