高级自动微分

在本教程中，您将了解 JAX 中自动微分 (autodiff) 的复杂应用，并更好地了解如何在 JAX 中轻松且强大地求导。

如果您还没有阅读过 自动微分 教程，请务必阅读该教程，以了解 JAX 自动微分的基础知识。

设置

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

key = random.key(0)

求梯度（第二部分）

高阶导数

JAX 的自动微分使得计算高阶导数变得容易，因为计算导数的函数本身是可微分的。因此，高阶导数就像堆叠变换一样容易。

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

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

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

多个变量的实值函数 \(f: \mathbb R^n\to\mathbb R\) 的 Hessian 可以与其梯度的 Jacobian 相对应。

JAX 提供了两个用于计算函数 Jacobian 的变换：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)

停止梯度

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

例如，考虑 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。我们通过向 jax.vmap() 添加 in_axes 来解决此问题，将 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()

6.77 ms ± 29.6 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
11.5 μs ± 24.1 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\) 是一个新的标量值函数，该函数将 \(f\) 在 \(x\) 处的梯度与向量 \(v\) 点积。请注意，您只对向量值参数的标量值函数求导，这正是您知道 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.10918739  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.10918739  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.09748875 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\)），这正是前向模式胜出的地方。

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

雅可比矩阵-向量积（JVP，又名：前向模式自动微分）

JAX 包括前向模式和反向模式自动微分的高效通用实现。我们熟悉的 jax.grad() 函数基于反向模式构建，但是为了解释两种模式之间的差异以及每种模式何时有用，您需要一些数学背景知识。

数学中的 JVP

在数学上，给定一个函数 \(f : \mathbb{R}^n \to \mathbb{R}^m\)，\(f\) 在输入点 \(x \in \mathbb{R}^n\) 处求值的雅可比矩阵，记为 \(\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 代码中的 JVP

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

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。

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

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

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

如果您正在机器学习中进行基于梯度的优化，您可能想要最小化从 \(\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
6.01 ms ± 101 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over forward
14.1 ms ± 9.44 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over reverse
20.4 ms ± 13.5 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Naive full Hessian materialization
54.9 ms ± 1.47 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
182 ms ± 336 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

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

/tmp/ipykernel_622/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
242 ms ± 244 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

Vmapped Jacobian-Matrix product
3.04 ms ± 101 μ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 雅可比矩阵具有复平面中缩放和旋转矩阵的特殊形式，即单个复数在乘法下的作用。）并且我们可以使用对带有 1.0 的余向量的 vjp 的单个调用来揭示该复数。

因为这仅适用于全纯函数，所以要使用此技巧，我们需要向 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.7534186  +0.j       , -3.0509028 -10.940544j ,
         5.9896846  +3.5423026j],
       [-3.0509028 +10.940544j , -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 的信息，请参阅关于添加原语的笔记本。

简而言之：使用 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

简而言之：使用 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 b
    _:f32[] = log c
    d:f32[] = div 1.0 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 : \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

@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 0x7f7c78753460>]

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

示例：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\)，因此我们可以认为存在一个由方程 \(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

我们也可以 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 的《非线性规划》第 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\) 展开后，您可以得出结论，只需要评估 \(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 来计算原始输出。在高阶微分的上下文中，只有当规则调用原始 f 来计算原始输出时，微分变换的每次应用才会使用自定义 JVP 规则。（这代表了一种根本的权衡，我们不能在我们的规则中使用来自 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 也存在类似的情况。

使用 nondiff_argnums 的 jax.custom_jvp

使用 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

使用 nondiff_argnums 的 jax.custom_vjp

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 用于整数数据类型的数组参数。

下一步

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

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

• 自定义 VJP 和 JVP

• 固定点的高效导数

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

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

• 对自定义数据类型求导

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

• 通过雅可比预先累积来优化 VJP