高级自动微分

在本教程中，您将学习 JAX 中自动微分 (autodiff) 的复杂应用，并更好地理解 JAX 中求导既简单又强大。

如果您还没有这样做，请务必查看自动微分教程，以复习 JAX autodiff 的基础知识。

设置

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

key = random.key(0)

求梯度（第 2 部分）

高阶导数

JAX 的 autodiff 可以轻松计算高阶导数，因为计算导数的函数本身是可微的。因此，高阶导数与堆叠变换一样容易。

单变量情况已在自动微分教程中介绍，其中示例展示了如何使用jax.grad()计算\(f(x) = x^3 + 2x^2 - 3x + 1\)的导数。

在多变量情况下，高阶导数更复杂。函数的二阶导数由其Hessian 矩阵表示，定义如下

\[(\mathbf{H}f)_{i,j} = \frac{\partial^2 f}{\partial_i\partial_j}.\]

多个变量的实值函数的 Hessian 矩阵，\(f: \mathbb R^n\to\mathbb R\)，可以与其梯度的雅可比矩阵识别。

JAX 提供了两种变换来计算函数的雅可比矩阵，jax.jacfwd()和jax.jacrev()，分别对应于正向和反向模式自动微分。它们给出相同的答案，但在不同情况下，一种可能比另一种更有效 - 请参阅关于自动微分的视频。

def hessian(f):
  return jax.jacfwd(jax.grad(f))

让我们仔细检查一下点积\(f: \mathbf{x} \mapsto \mathbf{x} ^\top \mathbf{x}\)是否正确。

如果\(i=j\)，\(\frac{\partial^2 f}{\partial_i\partial_j}(\mathbf{x}) = 2\)。否则，\(\frac{\partial^2 f}{\partial_i\partial_j}(\mathbf{x}) = 0\)。

def f(x):
  return jnp.dot(x, x)

hessian(f)(jnp.array([1., 2., 3.]))

Array([[2., 0., 0.],
       [0., 2., 0.],
       [0., 0., 2.]], dtype=float32)

高阶优化

一些元学习技术，例如与模型无关的元学习（MAML），需要通过梯度更新进行微分。在其他框架中，这可能非常繁琐，但在 JAX 中，这要容易得多

def meta_loss_fn(params, data):
  """Computes the loss after one step of SGD."""
  grads = jax.grad(loss_fn)(params, data)
  return loss_fn(params - lr * grads, data)

meta_grads = jax.grad(meta_loss_fn)(params, data)

停止梯度

Autodiff 能够自动计算函数关于其输入的梯度。但是，有时您可能需要一些额外的控制：例如，您可能希望避免通过计算图的某些子集反向传播梯度。

例如，考虑 TD(0)（时间差分）强化学习更新。这用于学习从与环境交互的经验中估计环境中状态的值。假设状态\(s_{t-1}\)中的值估计\(v_{\theta}(s_{t-1}\))由线性函数参数化。

# Value function and initial parameters
value_fn = lambda theta, state: jnp.dot(theta, state)
theta = jnp.array([0.1, -0.1, 0.])

考虑从状态\(s_{t-1}\)到状态\(s_t\)的过渡，在此期间您观察到奖励\(r_t\)

# An example transition.
s_tm1 = jnp.array([1., 2., -1.])
r_t = jnp.array(1.)
s_t = jnp.array([2., 1., 0.])

网络参数的 TD(0) 更新为

\[ \Delta \theta = (r_t + v_{\theta}(s_t) - v_{\theta}(s_{t-1})) \nabla v_{\theta}(s_{t-1}) \]

此更新不是任何损失函数的梯度。

但是，它可以写成伪损失函数的梯度

\[ L(\theta) = - \frac{1}{2} [r_t + v_{\theta}(s_t) - v_{\theta}(s_{t-1})]^2 \]

如果目标\(r_t + v_{\theta}(s_t)\)对参数\(\theta\)的依赖性被忽略。

如何在 JAX 中实现这一点？如果您天真地编写伪损失，您会得到

def td_loss(theta, s_tm1, r_t, s_t):
  v_tm1 = value_fn(theta, s_tm1)
  target = r_t + value_fn(theta, s_t)
  return -0.5 * ((target - v_tm1) ** 2)

td_update = jax.grad(td_loss)
delta_theta = td_update(theta, s_tm1, r_t, s_t)

delta_theta

Array([-1.2,  1.2, -1.2], dtype=float32)

但是td_update将不计算 TD(0) 更新，因为梯度计算将包括target对\(\theta\)的依赖性。

您可以使用jax.lax.stop_gradient()强制 JAX 忽略目标对\(\theta\)的依赖性

def td_loss(theta, s_tm1, r_t, s_t):
  v_tm1 = value_fn(theta, s_tm1)
  target = r_t + value_fn(theta, s_t)
  return -0.5 * ((jax.lax.stop_gradient(target) - v_tm1) ** 2)

td_update = jax.grad(td_loss)
delta_theta = td_update(theta, s_tm1, r_t, s_t)

delta_theta

Array([ 1.2,  2.4, -1.2], dtype=float32)

这将target视为好像它不依赖于参数\(\theta\)，并计算参数的正确更新。

现在，让我们也使用原始 TD(0) 更新表达式计算\(\Delta \theta\)，以交叉检查我们的工作。您可能希望尝试使用jax.grad()和您到目前为止的知识来自己实现此操作。这是我们的解决方案

s_grad = jax.grad(value_fn)(theta, s_tm1)
delta_theta_original_calculation = (r_t + value_fn(theta, s_t) - value_fn(theta, s_tm1)) * s_grad

delta_theta_original_calculation # [1.2, 2.4, -1.2], same as `delta_theta`

Array([ 1.2,  2.4, -1.2], dtype=float32)

jax.lax.stop_gradient在其他设置中也可能很有用，例如，如果您希望来自某些损失的梯度仅影响神经网络的参数子集（例如，因为其他参数使用不同的损失进行训练）。

使用 stop_gradient 的直接估计器

直接估计器是一种技巧，用于定义通常不可微分的函数的“梯度”。给定一个非微分函数\(f : \mathbb{R}^n \to \mathbb{R}^n\)，它用作我们希望找到梯度的较大函数的一部分，我们只需在反向传递期间假装\(f\)是恒等函数。这可以使用jax.lax.stop_gradient整齐地实现

def f(x):
  return jnp.round(x)  # non-differentiable

def straight_through_f(x):
  # Create an exactly-zero expression with Sterbenz lemma that has
  # an exactly-one gradient.
  zero = x - jax.lax.stop_gradient(x)
  return zero + jax.lax.stop_gradient(f(x))

print("f(x): ", f(3.2))
print("straight_through_f(x):", straight_through_f(3.2))

print("grad(f)(x):", jax.grad(f)(3.2))
print("grad(straight_through_f)(x):", jax.grad(straight_through_f)(3.2))

f(x):  3.0
straight_through_f(x): 3.0
grad(f)(x): 0.0
grad(straight_through_f)(x): 1.0

逐个示例的梯度

虽然大多数 ML 系统从批量数据计算梯度和更新，但出于计算效率和/或方差减少的原因，有时需要访问与批次中每个特定样本关联的梯度/更新。

例如，需要这样做才能根据梯度幅度对数据进行优先级排序，或在逐个样本的基础上应用裁剪/归一化。

在许多框架（PyTorch、TF、Theano）中，计算逐个示例的梯度通常并非易事，因为库直接累积批量的梯度。简单的解决方法（例如，计算每个示例的单独损失，然后聚合生成的梯度）通常效率非常低。

在 JAX 中，您可以轻松但高效地定义代码来计算每个样本的梯度。

只需将jax.jit()、jax.vmap()和jax.grad()变换组合在一起

perex_grads = jax.jit(jax.vmap(jax.grad(td_loss), in_axes=(None, 0, 0, 0)))

# Test it:
batched_s_tm1 = jnp.stack([s_tm1, s_tm1])
batched_r_t = jnp.stack([r_t, r_t])
batched_s_t = jnp.stack([s_t, s_t])

perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

Array([[ 1.2,  2.4, -1.2],
       [ 1.2,  2.4, -1.2]], dtype=float32)

让我们一次完成一个变换。

首先，您将jax.grad()应用于td_loss以获得一个函数，该函数计算单个（未批量）输入的参数的损失梯度

dtdloss_dtheta = jax.grad(td_loss)

dtdloss_dtheta(theta, s_tm1, r_t, s_t)

Array([ 1.2,  2.4, -1.2], dtype=float32)

此函数计算上述数组的一行。

然后，您使用jax.vmap()向量化此函数。这会将批量维度添加到所有输入和输出。现在，给定一批输入，您会生成一批输出 - 批次中的每个输出对应于输入批次的相应成员的梯度。

almost_perex_grads = jax.vmap(dtdloss_dtheta)

batched_theta = jnp.stack([theta, theta])
almost_perex_grads(batched_theta, batched_s_tm1, batched_r_t, batched_s_t)

Array([[ 1.2,  2.4, -1.2],
       [ 1.2,  2.4, -1.2]], dtype=float32)

这并不是我们想要的，因为我们必须手动为此函数提供一批theta，而我们实际上想要使用单个theta。我们通过将in_axes添加到jax.vmap()来解决此问题，将 theta 指定为None，将其他参数指定为0。这会使生成的函数仅将额外的轴添加到其他参数，而保持theta未批量处理，正如我们所希望的那样

inefficient_perex_grads = jax.vmap(dtdloss_dtheta, in_axes=(None, 0, 0, 0))

inefficient_perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

Array([[ 1.2,  2.4, -1.2],
       [ 1.2,  2.4, -1.2]], dtype=float32)

这确实实现了我们想要的功能，但比它必须的速度慢。现在，您可以将整个内容包装在jax.jit()中，以获得相同函数的已编译、高效版本

perex_grads = jax.jit(inefficient_perex_grads)

perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

Array([[ 1.2,  2.4, -1.2],
       [ 1.2,  2.4, -1.2]], dtype=float32)

%timeit inefficient_perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t).block_until_ready()
%timeit perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t).block_until_ready()

5.49 ms ± 58.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
6.72 μs ± 29.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

使用 jax.grad 的 jax.grad 的 Hessian-向量积

您可以使用高阶jax.grad()做的一件事是构建 Hessian-向量积函数。（稍后您将编写一个更高效的实现，它混合使用正向和反向模式，但这个实现将使用纯反向模式。）

Hessian-向量积函数在截断牛顿共轭梯度算法中很有用，用于最小化平滑凸函数，或用于研究神经网络训练目标的曲率（例如，1、2、3、4）。

对于具有连续二阶导数的标量值函数\(f : \mathbb{R}^n \to \mathbb{R}\)（因此 Hessian 矩阵是对称的），点\(x \in \mathbb{R}^n\)处的 Hessian 写为\(\partial^2 f(x)\)。然后，Hessian-向量积函数能够评估

\(\qquad v \mapsto \partial^2 f(x) \cdot v\)

对于任何\(v \in \mathbb{R}^n\)。

诀窍是不实例化完整的 Hessian 矩阵：如果\(n\)很大，在神经网络的上下文中可能数百万或数十亿，那么可能无法存储。

幸运的是，jax.grad()已经为我们提供了一种编写高效 Hessian-向量积函数的方法。您只需要使用恒等式

\(\qquad \partial^2 f (x) v = \partial [x \mapsto \partial f(x) \cdot v] = \partial g(x)\),

其中\(g(x) = \partial f(x) \cdot v\)是一个新的标量值函数，它用向量\(v\)点乘\(f\)在\(x\)处的梯度。请注意，您只对向量值参数的标量值函数进行微分，这正是您知道jax.grad()高效的地方。

在 JAX 代码中，您可以直接编写

def hvp(f, x, v):
    return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)

此示例表明您可以自由使用词法闭包，JAX 永远不会被打扰或困惑。

一旦您学会了如何计算密集 Hessian 矩阵，您将在几个单元格中检查此实现。您还将编写一个更好的版本，它同时使用正向模式和反向模式。

使用 jax.jacfwd 和 jax.jacrev 的雅可比矩阵和 Hessian 矩阵

您可以使用jax.jacfwd()和jax.jacrev()函数计算完整的雅可比矩阵

from jax import jacfwd, jacrev

# Define a sigmoid function.
def sigmoid(x):
    return 0.5 * (jnp.tanh(x / 2) + 1)

# Outputs probability of a label being true.
def predict(W, b, inputs):
    return sigmoid(jnp.dot(inputs, W) + b)

# Build a toy dataset.
inputs = jnp.array([[0.52, 1.12,  0.77],
                   [0.88, -1.08, 0.15],
                   [0.52, 0.06, -1.30],
                   [0.74, -2.49, 1.39]])

# Initialize random model coefficients
key, W_key, b_key = random.split(key, 3)
W = random.normal(W_key, (3,))
b = random.normal(b_key, ())

# Isolate the function from the weight matrix to the predictions
f = lambda W: predict(W, b, inputs)

J = jacfwd(f)(W)
print("jacfwd result, with shape", J.shape)
print(J)

J = jacrev(f)(W)
print("jacrev result, with shape", J.shape)
print(J)

jacfwd result, with shape (4, 3)
[[ 0.05069415  0.1091874   0.07506633]
 [ 0.14170025 -0.17390487  0.02415345]
 [ 0.12579198  0.01451446 -0.31447992]
 [ 0.00574409 -0.0193281   0.01078958]]
jacrev result, with shape (4, 3)
[[ 0.05069415  0.10918741  0.07506634]
 [ 0.14170025 -0.17390487  0.02415345]
 [ 0.12579198  0.01451446 -0.31447995]
 [ 0.00574409 -0.0193281   0.01078958]]

这两个函数计算相同的值（直到机器数值），但在它们的实现上有所不同：jax.jacfwd()使用正向模式自动微分，对于“高”雅可比矩阵（输出多于输入）更有效，而jax.jacrev()使用反向模式，对于“宽”雅可比矩阵（输入多于输出）更有效。对于接近正方形的矩阵，jax.jacfwd()可能比jax.jacrev()更有优势。

您还可以将jax.jacfwd()和jax.jacrev()与容器类型一起使用

def predict_dict(params, inputs):
    return predict(params['W'], params['b'], inputs)

J_dict = jacrev(predict_dict)({'W': W, 'b': b}, inputs)
for k, v in J_dict.items():
    print("Jacobian from {} to logits is".format(k))
    print(v)

Jacobian from W to logits is
[[ 0.05069415  0.10918741  0.07506634]
 [ 0.14170025 -0.17390487  0.02415345]
 [ 0.12579198  0.01451446 -0.31447995]
 [ 0.00574409 -0.0193281   0.01078958]]
Jacobian from b to logits is
[0.09748876 0.16102302 0.24190766 0.00776229]

有关正向和反向模式的更多详细信息，以及如何尽可能高效地实现jax.jacfwd()和jax.jacrev()，请继续阅读！

使用这两个函数的组合，我们可以计算密集 Hessian 矩阵

def hessian(f):
    return jacfwd(jacrev(f))

H = hessian(f)(W)
print("hessian, with shape", H.shape)
print(H)

hessian, with shape (4, 3, 3)
[[[ 0.02058932  0.04434624  0.03048803]
  [ 0.04434623  0.09551499  0.06566654]
  [ 0.03048803  0.06566655  0.04514575]]

 [[-0.0743913   0.09129842 -0.01268033]
  [ 0.09129842 -0.11204806  0.01556223]
  [-0.01268034  0.01556223 -0.00216142]]

 [[ 0.01176856  0.00135791 -0.02942139]
  [ 0.00135791  0.00015668 -0.00339478]
  [-0.0294214  -0.00339478  0.07355348]]

 [[-0.00418412  0.014079   -0.00785936]
  [ 0.014079   -0.04737393  0.02644569]
  [-0.00785936  0.02644569 -0.01476286]]]

这种形状是有道理的：如果您从函数\(f : \mathbb{R}^n \to \mathbb{R}^m\)开始，那么在点\(x \in \mathbb{R}^n\)处，您希望得到以下形状

• \(f(x) \in \mathbb{R}^m\)，\(f\)在\(x\)处的值，

• \(\partial f(x) \in \mathbb{R}^{m \times n}\)，\(x\)处的雅可比矩阵，

• \(\partial^2 f(x) \in \mathbb{R}^{m \times n \times n}\)，\(x\)处的 Hessian，

等等。

要实现hessian，您可以使用jacfwd(jacrev(f))或jacrev(jacfwd(f))或这两个函数的任何其他组合。但正向覆盖反向通常是最有效的。这是因为在内部雅可比计算中，我们经常对宽雅可比函数进行微分（可能像损失函数\(f : \mathbb{R}^n \to \mathbb{R}\)），而在外部雅可比计算中，我们对具有正方形雅可比矩阵的函数进行微分（因为\(\nabla f : \mathbb{R}^n \to \mathbb{R}^n\)），这正是正向模式胜出的地方。

它是如何制作的：两个基本的自动微分函数

雅可比-向量积（JVPs，又名正向模式自动微分）

JAX 包括正向和反向模式自动微分的高效且通用的实现。熟悉的jax.grad()函数建立在反向模式之上，但要解释这两种模式之间的区别，以及每种模式何时有用，您需要一些数学背景知识。

JVPs 的数学原理

从数学上讲，给定一个函数\(f : \mathbb{R}^n \to \mathbb{R}^m\)，在输入点\(x \in \mathbb{R}^n\)处评估的\(f\)的雅可比矩阵，表示为\(\partial f(x)\)，通常被认为是\(\mathbb{R}^m \times \mathbb{R}^n\)中的矩阵

\(\qquad \partial f(x) \in \mathbb{R}^{m \times n}\).

但您也可以将\(\partial f(x)\)视为线性映射，它将\(f\)域在点\(x\)处的切线空间（这只是\(\mathbb{R}^n\)的另一个副本）映射到\(f\)的陪域在点\(f(x)\)处的切线空间（\(\mathbb{R}^m\)的副本）

\(\qquad \partial f(x) : \mathbb{R}^n \to \mathbb{R}^m\).

此映射称为\(f\)在\(x\)处的推前映射。雅可比矩阵只是此线性映射在标准基上的矩阵。

如果您不提交到某个特定的输入点\(x\)，那么您可以将函数\(\partial f\)视为首先获取一个输入点，然后返回该输入点处的雅可比线性映射

\(\qquad \partial f : \mathbb{R}^n \to \mathbb{R}^n \to \mathbb{R}^m\).

特别是，您可以取消柯里化，以便给定输入点\(x \in \mathbb{R}^n\)和一个切向量\(v \in \mathbb{R}^n\)，您会得到\(\mathbb{R}^m\)中的输出切向量。我们将从\((x, v)\)对到输出切向量的映射称为雅可比-向量积，并将其写为

\(\qquad (x, v) \mapsto \partial f(x) v\)

JAX 代码中的 JVPs

回到 Python 代码中，JAX 的jax.jvp()函数对此变换进行建模。给定一个评估\(f\)的 Python 函数，JAX 的jax.jvp()是一种获取评估\((x, v) \mapsto (f(x), \partial f(x) v)\)的 Python 函数的方法。

from jax import jvp

# Isolate the function from the weight matrix to the predictions
f = lambda W: predict(W, b, inputs)

key, subkey = random.split(key)
v = random.normal(subkey, W.shape)

# Push forward the vector `v` along `f` evaluated at `W`
y, u = jvp(f, (W,), (v,))

就类似于 Haskell 的类型签名而言，您可以编写

jvp :: (a -> b) -> a -> T a -> (b, T b)

其中T a用于表示a的切线空间的类型。

换句话说，jvp将类型为a -> b的函数、类型为a的值和类型为T a的切向量值作为参数。它返回一对，包括类型为b的值和类型为T b的输出切向量。

jvp转换后的函数的评估方式与原始函数非常相似，但与类型为a的每个原始值配对，它沿着类型为T a的切线值进行推送。对于原始函数将应用的每个原始数值运算，jvp转换后的函数都会执行该原始函数的“JVP 规则”，该规则既可以在原始值上评估原始函数，也可以在该原始值上应用原始函数的 JVP。

该评估策略对计算复杂性有一些直接的影响。由于我们一边进行一边评估 JVPs，所以我们不需要存储任何东西供以后使用，因此内存成本与计算深度无关。此外，jvp转换后的函数的 FLOP 成本大约是仅评估函数成本的 3 倍（评估原始函数的一个单位工作量，例如 sin(x)；线性化的一个单位，例如 cos(x)；以及将线性化函数应用于向量的一个单位，例如 cos_x * v）。换句话说，对于固定的原始点 \(x\)，我们可以评估 \(v \mapsto \partial f(x) \cdot v\) 的边际成本与评估 \(f\) 的成本大致相同。

这种内存复杂度听起来非常引人注目！那么为什么我们在机器学习中不经常看到前向模式呢？

要回答这个问题，首先考虑一下如何使用 JVP 构建完整的雅可比矩阵。如果我们将 JVP 应用于 one-hot 切向量，它会揭示雅可比矩阵的一列，对应于我们输入的非零条目。因此，我们可以一次构建一个完整的雅可比矩阵，而获取每一列的成本与一次函数评估的成本大致相同。这对于具有“高”雅可比矩阵的函数来说是有效的，但对于“宽”雅可比矩阵来说是低效的。

如果您在机器学习中进行基于梯度的优化，您可能希望最小化从 \(\mathbb{R}^n\) 中的参数到 \(\mathbb{R}\) 中的标量损失值的损失函数。这意味着此函数的雅可比矩阵是一个非常宽的矩阵：\(\partial f(x) \in \mathbb{R}^{1 \times n}\)，我们通常将其与梯度向量 \(\nabla f(x) \in \mathbb{R}^n\) 识别。一次构建一列该矩阵，每次调用花费与评估原始函数相似的 FLOP 数量，这似乎效率低下！特别是，对于训练神经网络，其中 \(f\) 是训练损失函数，并且 \(n\) 可以达到数百万或数十亿，这种方法将无法扩展。

要更好地处理此类函数，您只需要使用反向模式。

向量-雅可比矩阵乘积 (VJP，又名反向模式自动微分)

前向模式返回一个用于评估雅可比矩阵-向量乘积的函数，然后我们可以使用该函数一次构建雅可比矩阵的一列，而反向模式是一种返回用于评估向量-雅可比矩阵乘积（等效于雅可比矩阵-转置-向量乘积）的函数的方法，我们可以使用该函数一次构建雅可比矩阵的一行。

VJP 在数学中的表示

让我们再次考虑一个函数 \(f : \mathbb{R}^n \to \mathbb{R}^m\)。从我们对 JVP 的表示法开始，VJP 的表示法非常简单

\(\qquad (x, v) \mapsto v \partial f(x)\),

其中 \(v\) 是 \(f\) 在 \(x\) 处的余切空间的一个元素（与 \(\mathbb{R}^m\) 的另一个副本同构）。在严格的情况下，我们应该将 \(v\) 视为线性映射 \(v : \mathbb{R}^m \to \mathbb{R}\)，当我们写 \(v \partial f(x)\) 时，我们的意思是函数组合 \(v \circ \partial f(x)\)，其中类型有效是因为 \(\partial f(x) : \mathbb{R}^n \to \mathbb{R}^m\)。但在常见情况下，我们可以将 \(v\) 与 \(\mathbb{R}^m\) 中的向量识别，并且几乎可以互换使用这两者，就像我们有时会在“列向量”和“行向量”之间切换而没有太多评论一样。

通过这种识别，我们也可以将 VJP 的线性部分视为 JVP 线性部分的转置（或伴随共轭）

\(\qquad (x, v) \mapsto \partial f(x)^\mathsf{T} v\).

对于给定的点 \(x\)，我们可以将签名写为

\(\qquad \partial f(x)^\mathsf{T} : \mathbb{R}^m \to \mathbb{R}^n\).

余切空间上的相应映射通常称为 \(f\) 在 \(x\) 处的回拉。对于我们的目的来说，关键是它从看起来像 \(f\) 的输出的东西到看起来像 \(f\) 的输入的东西，就像我们可能从转置线性函数中期望的那样。

JAX 代码中的 VJP

从数学切换回 Python，JAX 函数 vjp 可以接受一个用于评估 \(f\) 的 Python 函数，并返回一个用于评估 VJP \((x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))\) 的 Python 函数。

from jax import vjp

# Isolate the function from the weight matrix to the predictions
f = lambda W: predict(W, b, inputs)

y, vjp_fun = vjp(f, W)

key, subkey = random.split(key)
u = random.normal(subkey, y.shape)

# Pull back the covector `u` along `f` evaluated at `W`
v = vjp_fun(u)

用 Haskell 类似类型签名来说，我们可以写

vjp :: (a -> b) -> a -> (b, CT b -> CT a)

其中我们使用 CT a 来表示 a 的余切空间的类型。用文字来说，vjp 接受一个类型为 a -> b 的函数和一个类型为 a 的点作为参数，并返回一个由类型为 b 的值和一个类型为 CT b -> CT a 的线性映射组成的对。

这很棒，因为它允许我们一次构建雅可比矩阵的一行，并且评估 \((x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))\) 的 FLOP 成本仅为评估 \(f\) 成本的三倍左右。特别是，如果我们想要函数 \(f : \mathbb{R}^n \to \mathbb{R}\) 的梯度，我们只需调用一次即可完成。这就是 jax.grad() 对于基于梯度的优化是有效的，即使对于数百万或数十亿个参数的神经网络训练损失函数等目标也是如此。

虽然 FLOP 很友好，但内存会随着计算的深度而扩展，因此存在成本。此外，该实现传统上比前向模式的实现更复杂，尽管 JAX 有一些技巧（这是未来笔记本的故事！）。

有关反向模式如何工作的更多信息，请查看 2017 年深度学习暑期学校的此教程视频。

带有 VJP 的向量值梯度

如果您有兴趣获取向量值梯度（如 tf.gradients）

def vgrad(f, x):
  y, vjp_fn = vjp(f, x)
  return vjp_fn(jnp.ones(y.shape))[0]

print(vgrad(lambda x: 3*x**2, jnp.ones((2, 2))))

[[6. 6.]
 [6. 6.]]

使用前向模式和反向模式的 Hessian-向量乘积

在上一节中，您仅使用反向模式实现了一个 Hessian-向量乘积函数（假设连续二阶导数）

def hvp(f, x, v):
    return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)

这很有效，但通过将前向模式与反向模式结合使用，您可以做得更好并节省一些内存。

在数学上，给定一个要微分的函数 \(f : \mathbb{R}^n \to \mathbb{R}\)、一个用于线性化函数的点 \(x \in \mathbb{R}^n\) 和一个向量 \(v \in \mathbb{R}^n\)，我们想要的 Hessian-向量乘积函数是

\((x, v) \mapsto \partial^2 f(x) v\)

考虑辅助函数 \(g : \mathbb{R}^n \to \mathbb{R}^n\)，它被定义为 \(f\) 的导数（或梯度），即 \(g(x) = \partial f(x)\)。您所需要的只是它的 JVP，因为这将给我们

\((x, v) \mapsto \partial g(x) v = \partial^2 f(x) v\).

我们可以几乎直接将其翻译成代码

# forward-over-reverse
def hvp(f, primals, tangents):
  return jvp(grad(f), primals, tangents)[1]

更好的是，由于您不必直接调用 jnp.dot()，因此此 hvp 函数可以处理任何形状的数组和任意容器类型（例如存储为嵌套列表/字典/元组的向量），甚至不依赖于 jax.numpy。

这是一个如何使用它的示例

def f(X):
  return jnp.sum(jnp.tanh(X)**2)

key, subkey1, subkey2 = random.split(key, 3)
X = random.normal(subkey1, (30, 40))
V = random.normal(subkey2, (30, 40))

ans1 = hvp(f, (X,), (V,))
ans2 = jnp.tensordot(hessian(f)(X), V, 2)

print(jnp.allclose(ans1, ans2, 1e-4, 1e-4))

True

您可以考虑编写此代码的另一种方式是使用反向模式在前向模式之上

# Reverse-over-forward
def hvp_revfwd(f, primals, tangents):
  g = lambda primals: jvp(f, primals, tangents)[1]
  return grad(g)(primals)

不过，这还不太好，因为前向模式的开销小于反向模式，并且由于此处的外部微分算子必须区分比内部微分更大的计算，因此将前向模式保持在外部效果最佳

# Reverse-over-reverse, only works for single arguments
def hvp_revrev(f, primals, tangents):
  x, = primals
  v, = tangents
  return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)


print("Forward over reverse")
%timeit -n10 -r3 hvp(f, (X,), (V,))
print("Reverse over forward")
%timeit -n10 -r3 hvp_revfwd(f, (X,), (V,))
print("Reverse over reverse")
%timeit -n10 -r3 hvp_revrev(f, (X,), (V,))

print("Naive full Hessian materialization")
%timeit -n10 -r3 jnp.tensordot(hessian(f)(X), V, 2)

Forward over reverse
4.59 ms ± 135 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over forward
11.9 ms ± 7.17 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over reverse
15.3 ms ± 6.3 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Naive full Hessian materialization
47.6 ms ± 1.74 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)

组合 VJP、JVP 和 jax.vmap

雅可比矩阵-矩阵和矩阵-雅可比矩阵乘积

现在您有了 jax.jvp() 和 jax.vjp() 转换，它们为您提供了每次推送或拉回单个向量的函数，您可以使用 JAX 的 jax.vmap() 转换来一次推送和拉回整个基。特别是，您可以使用它来编写快速矩阵-雅可比矩阵和雅可比矩阵-矩阵乘积

# Isolate the function from the weight matrix to the predictions
f = lambda W: predict(W, b, inputs)

# Pull back the covectors `m_i` along `f`, evaluated at `W`, for all `i`.
# First, use a list comprehension to loop over rows in the matrix M.
def loop_mjp(f, x, M):
    y, vjp_fun = vjp(f, x)
    return jnp.vstack([vjp_fun(mi) for mi in M])

# Now, use vmap to build a computation that does a single fast matrix-matrix
# multiply, rather than an outer loop over vector-matrix multiplies.
def vmap_mjp(f, x, M):
    y, vjp_fun = vjp(f, x)
    outs, = vmap(vjp_fun)(M)
    return outs

key = random.key(0)
num_covecs = 128
U = random.normal(key, (num_covecs,) + y.shape)

loop_vs = loop_mjp(f, W, M=U)
print('Non-vmapped Matrix-Jacobian product')
%timeit -n10 -r3 loop_mjp(f, W, M=U)

print('\nVmapped Matrix-Jacobian product')
vmap_vs = vmap_mjp(f, W, M=U)
%timeit -n10 -r3 vmap_mjp(f, W, M=U)

assert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Matrix-Jacobian Products should be identical'

Non-vmapped Matrix-Jacobian product
61.5 ms ± 58.7 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

Vmapped Matrix-Jacobian product
4.46 ms ± 58.2 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

/tmp/ipykernel_1323/3769736790.py:8: DeprecationWarning: vstack requires ndarray or scalar arguments, got <class 'tuple'> at position 0. In a future JAX release this will be an error.
  return jnp.vstack([vjp_fun(mi) for mi in M])

def loop_jmp(f, W, M):
    # jvp immediately returns the primal and tangent values as a tuple,
    # so we'll compute and select the tangents in a list comprehension
    return jnp.vstack([jvp(f, (W,), (mi,))[1] for mi in M])

def vmap_jmp(f, W, M):
    _jvp = lambda s: jvp(f, (W,), (s,))[1]
    return vmap(_jvp)(M)

num_vecs = 128
S = random.normal(key, (num_vecs,) + W.shape)

loop_vs = loop_jmp(f, W, M=S)
print('Non-vmapped Jacobian-Matrix product')
%timeit -n10 -r3 loop_jmp(f, W, M=S)
vmap_vs = vmap_jmp(f, W, M=S)
print('\nVmapped Jacobian-Matrix product')
%timeit -n10 -r3 vmap_jmp(f, W, M=S)

assert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Jacobian-Matrix products should be identical'

Non-vmapped Jacobian-Matrix product
164 ms ± 579 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

Vmapped Jacobian-Matrix product
2.29 ms ± 41.3 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

jax.jacfwd 和 jax.jacrev 的实现

现在我们已经看到了快速雅可比矩阵-矩阵和矩阵-雅可比矩阵乘积，不难猜测如何编写 jax.jacfwd() 和 jax.jacrev()。我们只是使用相同的技术来一次推送或拉回整个标准基（与单位矩阵同构）。

from jax import jacrev as builtin_jacrev

def our_jacrev(f):
    def jacfun(x):
        y, vjp_fun = vjp(f, x)
        # Use vmap to do a matrix-Jacobian product.
        # Here, the matrix is the Euclidean basis, so we get all
        # entries in the Jacobian at once.
        J, = vmap(vjp_fun, in_axes=0)(jnp.eye(len(y)))
        return J
    return jacfun

assert jnp.allclose(builtin_jacrev(f)(W), our_jacrev(f)(W)), 'Incorrect reverse-mode Jacobian results!'

from jax import jacfwd as builtin_jacfwd

def our_jacfwd(f):
    def jacfun(x):
        _jvp = lambda s: jvp(f, (x,), (s,))[1]
        Jt = vmap(_jvp, in_axes=1)(jnp.eye(len(x)))
        return jnp.transpose(Jt)
    return jacfun

assert jnp.allclose(builtin_jacfwd(f)(W), our_jacfwd(f)(W)), 'Incorrect forward-mode Jacobian results!'

有趣的是，Autograd 库无法做到这一点。Autograd 中反向模式 jacobian 的 实现必须使用外部循环 map 一次拉回一个向量。一次通过计算推送一个向量的效率远低于使用 jax.vmap() 将其全部批处理在一起。

Autograd 无法做的另一件事是 jax.jit()。有趣的是，无论您在要微分的函数中使用多少 Python 动态性，我们始终可以在计算的线性部分使用 jax.jit()。例如

def f(x):
    try:
        if x < 3:
            return 2 * x ** 3
        else:
            raise ValueError
    except ValueError:
        return jnp.pi * x

y, f_vjp = vjp(f, 4.)
print(jit(f_vjp)(1.))

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

复数和微分

JAX 在复数和微分方面非常出色。为了同时支持 全纯和非全纯微分，有助于从 JVP 和 VJP 的角度进行思考。

考虑一个复数到复数函数 \(f: \mathbb{C} \to \mathbb{C}\)，并将其与相应的函数 \(g: \mathbb{R}^2 \to \mathbb{R}^2\) 识别，

def f(z):
  x, y = jnp.real(z), jnp.imag(z)
  return u(x, y) + v(x, y) * 1j

def g(x, y):
  return (u(x, y), v(x, y))

也就是说，我们分解了 \(f(z) = u(x, y) + v(x, y) i\)，其中 \(z = x + y i\)，并将 \(\mathbb{C}\) 与 \(\mathbb{R}^2\) 识别以获得 \(g\)。

由于 \(g\) 仅涉及实数输入和输出，因此我们已经知道如何为它编写雅可比矩阵-向量乘积，例如给定一个切向量 \((c, d) \in \mathbb{R}^2\)，即

\(\begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \\ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} c \\ d \end{bmatrix}\).

要获得应用于切向量 \(c + di \in \mathbb{C}\) 的原始函数 \(f\) 的 JVP，我们只需使用相同的定义并将结果识别为另一个复数，

\(\partial f(x + y i)(c + d i) = \begin{matrix} \begin{bmatrix} 1 & i \end{bmatrix} \\ ~ \end{matrix} \begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \\ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} c \\ d \end{bmatrix}\).

这是我们对 \(\mathbb{C} \to \mathbb{C}\) 函数的 JVP 的定义！请注意，\(f\) 是否全纯无关紧要：JVP 是明确的。

这是一个检查

def check(seed):
  key = random.key(seed)

  # random coeffs for u and v
  key, subkey = random.split(key)
  a, b, c, d = random.uniform(subkey, (4,))

  def fun(z):
    x, y = jnp.real(z), jnp.imag(z)
    return u(x, y) + v(x, y) * 1j

  def u(x, y):
    return a * x + b * y

  def v(x, y):
    return c * x + d * y

  # primal point
  key, subkey = random.split(key)
  x, y = random.uniform(subkey, (2,))
  z = x + y * 1j

  # tangent vector
  key, subkey = random.split(key)
  c, d = random.uniform(subkey, (2,))
  z_dot = c + d * 1j

  # check jvp
  _, ans = jvp(fun, (z,), (z_dot,))
  expected = (grad(u, 0)(x, y) * c +
              grad(u, 1)(x, y) * d +
              grad(v, 0)(x, y) * c * 1j+
              grad(v, 1)(x, y) * d * 1j)
  print(jnp.allclose(ans, expected))

check(0)
check(1)
check(2)

True
True
True

VJP 怎么样？我们做了一些非常相似的事情：对于余切向量 \(c + di \in \mathbb{C}\)，我们将 \(f\) 的 VJP 定义为

\((c + di)^* \; \partial f(x + y i) = \begin{matrix} \begin{bmatrix} c & -d \end{bmatrix} \\ ~ \end{matrix} \begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \\ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} 1 \\ -i \end{bmatrix}\).

负数是怎么回事？它们只是为了处理复共轭，以及我们正在使用余向量的事实。

这是对 VJP 规则的检查

def check(seed):
  key = random.key(seed)

  # random coeffs for u and v
  key, subkey = random.split(key)
  a, b, c, d = random.uniform(subkey, (4,))

  def fun(z):
    x, y = jnp.real(z), jnp.imag(z)
    return u(x, y) + v(x, y) * 1j

  def u(x, y):
    return a * x + b * y

  def v(x, y):
    return c * x + d * y

  # primal point
  key, subkey = random.split(key)
  x, y = random.uniform(subkey, (2,))
  z = x + y * 1j

  # cotangent vector
  key, subkey = random.split(key)
  c, d = random.uniform(subkey, (2,))
  z_bar = jnp.array(c + d * 1j)  # for dtype control

  # check vjp
  _, fun_vjp = vjp(fun, z)
  ans, = fun_vjp(z_bar)
  expected = (grad(u, 0)(x, y) * c +
              grad(v, 0)(x, y) * (-d) +
              grad(u, 1)(x, y) * c * (-1j) +
              grad(v, 1)(x, y) * (-d) * (-1j))
  assert jnp.allclose(ans, expected, atol=1e-5, rtol=1e-5)

check(0)
check(1)
check(2)

像 jax.grad()、jax.jacfwd() 和 jax.jacrev() 这样的便利包装器怎么样？

对于 \(\mathbb{R} \to \mathbb{R}\) 函数，回想一下我们将 grad(f)(x) 定义为 vjp(f, x)[1](1.0)，这是可行的，因为将 VJP 应用于 1.0 值会显示梯度（即雅可比矩阵或导数）。我们可以对 \(\mathbb{C} \to \mathbb{R}\) 函数执行相同的操作：我们仍然可以使用 1.0 作为余切向量，我们只会得到一个复数结果，总结了完整的雅可比矩阵

def f(z):
  x, y = jnp.real(z), jnp.imag(z)
  return x**2 + y**2

z = 3. + 4j
grad(f)(z)

Array(6.-8.j, dtype=complex64)

对于一般的 \(\mathbb{C} \to \mathbb{C}\) 函数，雅可比矩阵具有 4 个实数值自由度（如上面的 2x2 雅可比矩阵中所示），因此我们无法希望在复数中表示所有这些自由度。但是我们可以为全纯函数做到这一点！全纯函数恰好是一个 \(\mathbb{C} \to \mathbb{C}\) 函数，其特殊性质是其导数可以表示为单个复数。（柯西-黎曼方程确保上述 2x2 雅可比矩阵具有复平面中缩放和旋转矩阵的特殊形式，即单个复数在乘法下的作用。）我们可以通过一次调用 vjp，并使用 1.0 的共向量来揭示该复数。

因为这仅适用于全纯函数，所以要使用此技巧，我们需要向 JAX 保证我们的函数是全纯的；否则，当 jax.grad() 用于复数输出函数时，JAX 会引发错误

def f(z):
  return jnp.sin(z)

z = 3. + 4j
grad(f, holomorphic=True)(z)

Array(-27.034946-3.8511534j, dtype=complex64, weak_type=True)

所有 holomorphic=True 承诺所做的只是在输出为复数值时禁用错误。当函数不是全纯函数时，我们仍然可以编写 holomorphic=True，但我们得到的答案不会表示完整的雅可比矩阵。相反，它将是函数的雅可比矩阵，我们在其中仅丢弃输出的虚部

def f(z):
  return jnp.conjugate(z)

z = 3. + 4j
grad(f, holomorphic=True)(z)  # f is not actually holomorphic!

Array(1.-0.j, dtype=complex64, weak_type=True)

对于 jax.grad() 在此处的工作方式，有一些有用的结果

1. 我们可以对全纯 \(\mathbb{C} \to \mathbb{C}\) 函数使用 jax.grad()。

2. 我们可以使用 jax.grad() 来优化 \(f : \mathbb{C} \to \mathbb{R}\) 函数，例如复数参数 x 的实数值损失函数，方法是沿着 grad(f)(x) 的共轭方向迈出一步。

3. 如果我们有一个 \(\mathbb{R} \to \mathbb{R}\) 函数，它恰好在内部使用一些复数值运算（其中一些必须是非全纯的，例如卷积中使用的 FFT），那么 jax.grad() 仍然有效，我们得到的结果与仅使用实数值的实现会给出的结果相同。

在任何情况下，JVP 和 VJP 始终是明确的。如果我们想计算非全纯 \(\mathbb{C} \to \mathbb{C}\) 函数的完整雅可比矩阵，我们可以使用 JVP 或 VJP 来完成！

您应该期望复数在 JAX 中的所有位置都有效。这是通过复数矩阵的 Cholesky 分解进行微分

A = jnp.array([[5.,    2.+3j,    5j],
              [2.-3j,   7.,  1.+7j],
              [-5j,  1.-7j,    12.]])

def f(X):
    L = jnp.linalg.cholesky(X)
    return jnp.sum((L - jnp.sin(L))**2)

grad(f, holomorphic=True)(A)

Array([[-0.75341904 +0.j       , -3.0509028 -10.940545j ,
         5.9896846  +3.5423026j],
       [-3.0509028 +10.940545j , -8.904491   +0.j       ,
        -5.1351523  -6.559373j ],
       [ 5.9896846  -3.5423026j, -5.1351523  +6.559373j ,
         0.01320427 +0.j       ]], dtype=complex64)

用于 JAX 可转换 Python 函数的自定义导数规则

在 JAX 中定义微分规则有两种方式：

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

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

此笔记本是关于 #1 的。要阅读有关 #2 的内容，请参阅 关于添加原语的笔记本。

TL;DR：使用 jax.custom_jvp() 的自定义 JVP

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

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

TL;DR：使用 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 可转换的

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:f32[] b
    _:f32[] = log c
    d:f32[] = div 1.0:f32[] c
    e:f32[] = mul d b
  in (e,) }

逐步执行 jaxpr 的评估方式，请注意最后一行将涉及将浮点数学四舍五入为 0 和 \(\infty\) 的值相乘，这不是一个好主意。也就是说，我们有效地评估了 lambda x: (1 / (1 + jnp.exp(x))) * jnp.exp(x) 对于大的 x，这有效地变成了 0. * jnp.inf。

与其生成如此大和如此小的值，希望浮点数不能总是提供的消除，我们宁愿将导数函数表示为更数值稳定的程序。特别是，我们可以编写一个程序，该程序更接近于评估相等的数学表达式 \(1 - \frac{1}{1 + e^x}\)，而没有看到任何消除。

这个问题很有趣，因为即使我们对 log1pexp 的定义已经是 JAX 可微分的（并且使用 jax.jit()、jax.vmap() 等进行转换），我们对将标准自动微分规则应用于组成 log1pexp 的原语并组成结果的结果不满意。相反，我们想指定如何对整个函数 log1pexp 进行微分，作为一个单元，从而更好地安排这些指数。

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

这是一个使用 jax.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(x) = \frac{x}{1 + \sqrt{x}}\) 的函数 \(f : \mathbb{R}_+ \to \mathbb{R}_+\)，其中我们采用 \(\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 规则！特别是，我们可以根据导数函数 \(x \mapsto \frac{\sqrt{x} + 2}{2(\sqrt{x} + 1)^2}\) 在 \(\mathbb{R}_+\) 上定义 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

@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

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

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

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

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 0x7594cfabfef0>]

示例：Python 调试

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

当试图追踪 nan 运行时错误的来源，或者仅仅是仔细检查正在传播的余切（梯度）值时，在反向传播中对应于原始计算中特定点的位置插入调试器可能很有用。你可以使用 jax.custom_vjp() 来实现。

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

示例：迭代实现的隐函数微分

这个例子在数学上相当深入！

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

例如，考虑这个 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\)，因此我们可以认为存在一个函数 \(a \mapsto x^*(a)\)，该函数由方程 \(x = f(a, x)\) 隐式定义。

我们可以使用 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

我们也可以 jax.vmap() 或 jax.jit() 该函数

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

[1.        1.4142135 1.7320509 2.       ]

由于 while_loop，我们无法应用反向模式自动微分，但事实证明我们无论如何也不想这样做：与其通过 fixed_point 的实现及其所有迭代进行微分，不如利用数学结构来做一些更节省内存的事情（在这种情况下也更节省 FLOP！）。我们可以改为使用隐函数定理 [Bertsekas 的 Nonlinear Programming, 2nd ed. 的 Prop 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\) 后，你可以得出结论，你只需要评估 \(f\) 在 \((a_0, x^*(a_0))\) 处的 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 不能闭包任何涉及微分的值。也就是说，你可能会注意到我们在 fixed_point 的参数列表中显式地保留了参数 a。对于此用例，请考虑使用低级原语 lax.custom_root，它允许在具有自定义求根函数的闭包变量中进行微分。

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

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

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

# 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)>

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 的效率就像我们手动编写规则一样

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

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)

你可以使用 Python 控制流与 jax.custom_jvp()

@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

# 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)

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 的输出必须是一个序列（例如，一个元组），其长度等于原始函数的参数数量。

所以多个参数像这样工作

@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 规则。如果评估了该函数，或者使用 jax.jit()，jax.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):
  import pdb; 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 容器（和其他 pytree）

你应该期望标准的 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

使用 jax.custom_jvp() 的可选 nondiff_argnums 参数来指示这些参数。这是一个带有 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，例如，具有整数 dtype 的参数。相反，nondiff_argnums 应该仅用于与 JAX 类型不对应的参数值（本质上是不对应于数组类型的参数值），例如 Python 可调用对象或字符串。如果 JAX 检测到 nondiff_argnums 指示的参数包含 JAX Tracer，则会引发错误。上面的 clip_gradient 函数是一个不使用 nondiff_argnums 用于整数 dtype 数组参数的好例子。

下一步

还有其他自动微分技巧和功能的世界。本教程未涵盖但值得追求的主题包括

• 高斯-牛顿向量积，线性化一次

• 自定义 VJP 和 JVP

• 固定点处的有效导数

• 使用随机 Hessian-向量积估计 Hessian 的迹

• 仅使用反向模式自动微分的前向模式自动微分

• 计算相对于自定义数据类型的导数

• 检查点（用于有效反向模式的二项式检查点，而不是模型快照）

• 使用雅可比矩阵预累积优化 VJP