高级自动微分

高级自动微分#

在本教程中，您将学习 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)

求梯度（第 2 部分）#

高阶导数#

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 忽略 target 对 \(\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)

这还不是我们想要的，因为我们必须手动为此函数提供一批 thetas，而我们实际上只想使用单个 theta。我们通过将 in_axes 添加到 jax.vmap() 来解决这个问题，将 theta 指定为 None，并将其他 args 指定为 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()

3.55 ms ± 11.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.79 μs ± 2.33 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

使用 `jax.grad`-of-`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` 的 Jacobian 矩阵和 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)) 或这两个函数的任何其他组合。但前向模式 over 反向模式通常是最有效的。这是因为在内部雅可比矩阵计算中，我们通常对宽雅可比矩阵的函数进行微分（可能像损失函数 \(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\)，\(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)\).

特别地，您可以柯里化（uncurry）一下，以便给定输入点 \(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。

这种评估策略对计算复杂度有一些直接的影响。由于我们在进行过程中评估 JVP，因此我们不需要存储任何东西以供以后使用，因此内存成本与计算深度无关。此外，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\) 可能达到数百万或数十亿，这种方法就行不通了。

为了更好地处理像这样的函数，您只需要使用反向模式。

向量-雅可比积 (VJPs，又名：反向模式自动微分)#

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

VJPs 的数学原理#

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

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

其中 \(v\) 是 \(f\) 在 \(x\) 处的余切空间的元素（同构于 \(\mathbb{R}^m\) 的另一个副本）。当严格来说，我们应该将 \(v\) 视为线性映射 \(v : \mathbb{R}^m \to \mathbb{R}\)，当我们写 \(v^\mathsf{T} \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 代码中的 VJPs#

从数学切换回 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




您可能考虑编写此代码的另一种方式是使用反向模式 over 前向模式


# 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
2.47 ms ± 82.9 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over forward
The slowest run took 5.55 times longer than the fastest. This could mean that an intermediate result is being cached.
7.36 ms ± 6.27 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Reverse over reverse
The slowest run took 5.10 times longer than the fastest. This could mean that an intermediate result is being cached.
10.5 ms ± 8.59 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
Naive full Hessian materialization
28.8 ms ± 664 μs 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
54.2 ms ± 105 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

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


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

Vmapped Jacobian-Matrix product
1.51 ms ± 29.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() 在这里如何工作，有一些有用的结果

我们可以对全纯 \(\mathbb{C} \to \mathbb{C}\) 函数使用 jax.grad()。
我们可以使用 jax.grad() 来优化 \(f : \mathbb{C} \to \mathbb{R}\) 函数，例如复数参数 x 的实值损失函数，方法是沿 grad(f)(x) 的共轭方向步进。
如果我们有一个 \(\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 中定义微分规则有两种方法

使用 jax.custom_jvp() 和 jax.custom_vjp() 为已经是 JAX 可转换的 Python 函数定义自定义微分规则；以及
定义新的 core.Primitive 实例以及它们的所有转换规则，例如调用来自其他系统的函数，如求解器、模拟器或通用数值计算系统。

本 notebook 介绍了 #1。要阅读关于 #2 的内容，请参阅关于添加原语的 notebook。

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 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 微分（并使用 jax.jit()、jax.vmap() 等进行转换），但我们对将标准自动微分规则应用于构成 log1pexp 的原语并组合结果并不满意。相反，我们希望指定应如何对整个函数 log1pexp 进行微分，作为一个单元，从而更好地排列这些指数。
这是 Python 函数的自定义导数规则的一个应用，这些函数已经是 JAX 可转换的：指定应如何对复合函数进行微分，同时仍然将其原始 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 0x7915f423a0b0>]







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







示例：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 的《非线性规划》第二版命题 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 的 VJP，方法是调用 fixed_point！此外，在展开 \(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 容器（如列表、元组、namedtuple 和字典）以及这些容器的嵌套版本都可以正常工作。一般来说，只要 pytree 的结构根据类型约束保持一致，则任何 pytree 都是允许的。
这是一个使用 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。
对于数组值参数，例如整数 dtype 的参数，你不需要使用 nondiff_argnums。相反，nondiff_argnums 应该只用于不对应于 JAX 类型（本质上不对应于数组类型）的参数值，例如 Python 可调用对象或字符串。如果 JAX 检测到由 nondiff_argnums 指示的参数包含 JAX Tracer，则会引发错误。上面的 clip_gradient 函数是不对整数 dtype 数组参数使用 nondiff_argnums 的一个很好的例子。





接下来步骤#
还有很多其他的自动微分技巧和功能。本教程中未涵盖但值得探索的主题包括

高斯-牛顿向量积，线性化一次
自定义 VJP 和 JVP
定点处的有效导数
使用随机 Hessian-向量积估计 Hessian 的迹
仅使用反向模式自动微分的前向模式自动微分
对自定义数据类型求导
检查点（用于高效反向模式的二项式检查点，而不是模型快照）
使用 Jacobian 预先累积优化 VJP


              

              
              
              
              
                
                  

       上一页
 使用 JIT 的控制流和逻辑运算符
 
    下一页
 外部回调


            
            
              
                

  
  
     目录
  
    
设置
求梯度（第 2 部分）
高阶导数


高阶优化
停止梯度
使用 stop_gradient 的直通估计器
每个样本的梯度
使用 jax.grad-of-jax.grad 的 Hessian 向量积
使用 jax.jacfwd 和 jax.jacrev 的 Jacobian 矩阵和 Hessian 矩阵


如何实现：两个基础的自动微分函数
Jacobian 矩阵-向量积（JVPs，又名正向模式自动微分）
数学中的 JVP
JAX 代码中的 JVP


向量-Jacobian 矩阵积（VJPs，又名反向模式自动微分）
数学中的 VJP
JAX 代码中的 VJP


使用 VJP 的向量值梯度
使用正向和反向模式的 Hessian 向量积


组合 VJP、JVP 和 jax.vmap
Jacobian 矩阵-矩阵积和矩阵-Jacobian 矩阵积
jax.jacfwd 和 jax.jacrev 的实现


复数和微分
JAX 可变换 Python 函数的自定义导数规则
TL;DR：使用 jax.custom_jvp() 的自定义 JVP
TL;DR：使用 jax.custom_vjp 的自定义 VJP
示例问题
示例：数值稳定性
示例：强制微分约定
示例：梯度裁剪
示例：Python 调试
示例：迭代实现的隐函数微分


jax.custom_jvp 和 jax.custom_vjp API 的基本用法
使用 jax.custom_jvp 定义正向模式（以及间接的反向模式）规则
使用 jax.custom_vjp 定义自定义的仅反向模式规则


更多功能和细节
使用 list / tuple / dict 容器（和其他 pytrees）
处理不可微参数
带有 nondiff_argnums 的 jax.custom_jvp
带有 nondiff_argnums 的 jax.custom_vjp






下一步


          
            

  
  
    
作者：JAX 团队

  
  
  
    

  © 版权所有 2024，JAX 团队。

高级自动微分

目录

高级自动微分#

设置#

求梯度（第 2 部分）#

高阶导数#

高阶优化#

停止梯度#

使用 stop_gradient 的直通估计器#

每个样本的梯度#

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

使用 jax.jacfwd 和 jax.jacrev 的 Jacobian 矩阵和 Hessian 矩阵#

它是如何实现的：两个基础的自动微分函数#

雅可比-向量积 (JVPs，又名：前向模式自动微分)#

JVPs 的数学原理#

JAX 代码中的 JVPs#

向量-雅可比积 (VJPs，又名：反向模式自动微分)#

VJPs 的数学原理#

JAX 代码中的 VJPs#

使用 VJP 的向量值梯度#

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

组合 VJP、JVP 和 jax.vmap#

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

jax.jacfwd 和 jax.jacrev 的实现#

复数和微分#

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

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

TL;DR: 使用 jax.custom_vjp 的自定义 VJP#

示例问题#

示例：数值稳定性#

示例：强制微分约定#

示例：梯度裁剪#

示例：Python 调试#

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

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

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

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

更多功能和详细信息#

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

处理不可微参数#

带有 nondiff_argnums 的 jax.custom_jvp#

带有 nondiff_argnums 的 jax.custom_vjp#

接下来步骤#

使用 `stop_gradient` 的直通估计器#

使用 `jax.grad`-of-`jax.grad` 的 Hessian 向量积#

使用 `jax.jacfwd` 和 `jax.jacrev` 的 Jacobian 矩阵和 Hessian 矩阵#

组合 VJP、JVP 和 `jax.vmap`#

`jax.jacfwd` 和 `jax.jacrev` 的实现#

TL;DR: 使用 `jax.custom_jvp()` 的自定义 JVP#

TL;DR: 使用 `jax.custom_vjp` 的自定义 VJP#

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

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

使用 `jax.custom_vjp` 定义自定义仅反向模式规则#

使用 `list` / `tuple` / `dict` 容器（和其他 pytree）#

带有 `nondiff_argnums` 的 `jax.custom_jvp`#

带有 `nondiff_argnums` 的 `jax.custom_vjp`#