進階自動微分

在本教學中，您將學習 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}) \]

此更新不是任何損失函式的梯度。

但是，如果忽略目標 \(r_t + v_{\theta}(s_t)\) 對參數 \(\theta\) 的依賴性，則可以將其寫成偽損失函式的梯度

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

如果忽略目標 \(r_t + v_{\theta}(s_t)\) 對參數 \(\theta\) 的依賴性。

如何在 JAX 中實作這個？如果您天真地編寫偽損失，您會得到

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

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

delta_theta

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

但 td_update 將不會計算 TD(0) 更新，因為梯度計算將包含 target 對 \(\theta\) 的依賴性。

您可以使用 jax.lax.stop_gradient() 強制 JAX 忽略目標對 \(\theta\) 的依賴性

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

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

delta_theta

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

這會將 target 視為不依賴參數 \(\theta\)，並計算參數的正確更新。

現在，讓我們也使用原始 TD(0) 更新運算式計算 \(\Delta \theta\)，以交叉檢查我們的工作。您可能希望嘗試使用 jax.grad() 和您目前所學的知識自行實作。這是我們的解決方案

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

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

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

jax.lax.stop_gradient 在其他設定中也可能很有用，例如，如果您希望來自某些損失的梯度僅影響神經網路參數的子集（例如，因為其他參數是使用不同的損失訓練的）。

使用 stop_gradient 的直接傳遞估算器

直接傳遞估算器是一種技巧，用於定義非微分函式的「梯度」。給定一個非微分函式 \(f : \mathbb{R}^n \to \mathbb{R}^n\)，它被用作我們希望找到梯度的較大函式的一部分，我們只需在反向傳遞期間假裝 \(f\) 是恆等函式。這可以使用 jax.lax.stop_gradient 簡潔地實作

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

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

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

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

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

逐範例梯度

雖然大多數 ML 系統從批次資料計算梯度和更新，但基於計算效率和/或變異數縮減的原因，有時需要存取與批次中每個特定樣本相關聯的梯度/更新。

例如，這是根據梯度幅度優先排序資料，或在逐樣本基礎上應用裁剪/正規化所必需的。

在許多框架（PyTorch、TF、Theano）中，計算逐範例梯度通常並非易事，因為函式庫直接累計批次中的梯度。天真的解決方法，例如計算每個範例的單獨損失，然後聚合產生的梯度，通常非常低效。

在 JAX 中，您可以定義程式碼以簡單而有效的方式計算逐樣本梯度。

只需將 jax.jit()、jax.vmap() 和 jax.grad() 轉換組合在一起

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

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

perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

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

讓我們一次完成一個轉換。

首先，將 jax.grad() 應用於 td_loss，以取得一個函式，該函式計算單個（未批次化）輸入的損失相對於參數的梯度

dtdloss_dtheta = jax.grad(td_loss)

dtdloss_dtheta(theta, s_tm1, r_t, s_t)

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

此函式計算上述陣列的一列。

然後，您可以使用 jax.vmap() 將此函式向量化。這會將批次維度新增至所有輸入和輸出。現在，給定一批輸入，您會產生一批輸出 — 批次中的每個輸出都對應於輸入批次中對應成員的梯度。

almost_perex_grads = jax.vmap(dtdloss_dtheta)

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

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

這並不是我們真正想要的，因為我們必須手動將一批 theta 饋送到此函式，而我們實際上想要使用單個 theta。我們透過將 in_axes 新增至 jax.vmap() 來修正此問題，將 theta 指定為 None，而其他引數指定為 0。這會使產生的函式僅將額外軸新增至其他引數，而將 theta 保持未批次化，正如我們所願

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

inefficient_perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

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

這確實符合我們的需求，但速度比應有的速度慢。現在，您將整個內容包裝在 jax.jit() 中，以取得相同函式的已編譯、高效版本

perex_grads = jax.jit(inefficient_perex_grads)

perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)

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

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

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）。

對於具有連續二階導數（因此 Hessian 矩陣是對稱的）的純量值函式 \(f : \mathbb{R}^n \to \mathbb{R}\)，點 \(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() 函式計算完整的 Jacobian 矩陣

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() 使用正向模式自動微分，這對於「高」Jacobian 矩陣（輸出多於輸入）更有效率，而 jax.jacrev() 使用反向模式，這對於「寬」Jacobian 矩陣（輸入多於輸出）更有效率。對於接近方形的矩陣，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\) 的 Jacobian 矩陣，

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

依此類推。

為了實作 hessian，你可以使用 jacfwd(jacrev(f)) 或 jacrev(jacfwd(f))，或這兩者的任何其他組合。但前向微分覆蓋反向微分通常是最有效率的。那是因為在內部的 Jacobian 計算中，我們常常對寬 Jacobian 函數（可能像是損失函數 \(f : \mathbb{R}^n \to \mathbb{R}\)）進行微分，而在外部的 Jacobian 計算中，我們是對具有方陣 Jacobian 的函數進行微分（因為 \(\nabla f : \mathbb{R}^n \to \mathbb{R}^n\)），這正是前向模式勝出的地方。

它是如何製成的：兩個基礎的自動微分函數

Jacobian-向量積 (JVPs，又稱前向模式自動微分)

JAX 包含了前向和反向模式自動微分的高效率且通用的實作。大家熟悉的 jax.grad() 函數是建立在反向模式之上，但為了說明這兩種模式之間的差異，以及何時每種模式會很有用，你需要一些數學背景知識。

JVPs 的數學表示

在數學上，給定一個函數 \(f : \mathbb{R}^n \to \mathbb{R}^m\)，\(f\) 在輸入點 \(x \in \mathbb{R}^n\) 評估的 Jacobian，表示為 \(\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\) 的 前推映射 (pushforward map)。 Jacobian 矩陣只是這個線性映射在標準基底上的矩陣表示。

如果你不限定於特定的輸入點 \(x\)，那麼你可以將函數 \(\partial f\) 視為首先接收一個輸入點，然後返回該輸入點的 Jacobian 線性映射

\(\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)\) 對到輸出切向量的映射稱為 Jacobian-向量積，並將其寫為

\(\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 的輸出切向量組成的pair。

jvp 轉換後的函數的評估方式很像原始函數，但它會將類型為 a 的每個原始值與類型為 T a 的切線值配對並沿著推送。對於原始函數將會應用的每個基本數值運算，jvp 轉換後的函數會對該基本運算執行「JVP 規則」，該規則既評估原始值上的基本運算，又在這些原始值上應用基本運算的 JVP。

這種評估策略對計算複雜度有一些直接的影響。由於我們在進行過程中評估 JVP，因此我們不需要為以後儲存任何東西，因此記憶體成本與計算深度無關。此外，jvp 轉換後的函數的 FLOP 成本約為僅評估函數成本的 3 倍（一個工作單位用於評估原始函數，例如 sin(x)；一個單位用於線性化，如 cos(x)；一個單位用於將線性化函數應用於向量，如 cos_x * v）。換句話說，對於固定的原始點 \(x\)，我們可以以大約與評估 \(f\) 相同的邊際成本來評估 \(v \mapsto \partial f(x) \cdot v\)。

這種記憶體複雜度聽起來相當吸引人！那麼為什麼我們在機器學習中不常看到前向模式呢？

為了回答這個問題，首先想想你如何使用 JVP 來建立完整的 Jacobian 矩陣。如果我們將 JVP 應用於 one-hot 切向量，它會揭示 Jacobian 矩陣的一列，對應於我們輸入的非零項。因此，我們可以一次建立一個完整的 Jacobian 矩陣列，而取得每一列的成本與一次函數評估的成本大致相同。這對於具有「高」 Jacobian 的函數來說是有效率的，但對於「寬」 Jacobian 來說是效率低的。

如果你在機器學習中進行基於梯度的最佳化，你可能想要最小化從 \(\mathbb{R}^n\) 中的參數到 \(\mathbb{R}\) 中的純量損失值的損失函數。這表示這個函數的 Jacobian 是一個非常寬的矩陣：\(\partial f(x) \in \mathbb{R}^{1 \times n}\)，我們通常將其識別為梯度向量 \(\nabla f(x) \in \mathbb{R}^n\)。一次建立一列這個矩陣，每次調用都花費與評估原始函數相似的 FLOP 數量，這肯定看起來效率很低！特別是對於訓練神經網路來說，其中 \(f\) 是訓練損失函數，而 \(n\) 可能達到數百萬或數十億，這種方法根本無法擴展。

為了對這類函數做得更好，你只需要使用反向模式。

向量-Jacobian 積 (VJPs，又稱反向模式自動微分)

前向模式給我們一個評估 Jacobian-向量積的函數，然後我們可以利用它一次建立一個 Jacobian 矩陣列，反向模式是一種取得評估向量-Jacobian 積（等效於 Jacobian-轉置-向量積）函數的方法，我們可以利用它一次建立一個 Jacobian 矩陣行。

VJPs 的數學表示

讓我們再次考慮一個函數 \(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\) 的 拉回 (pullback)。對於我們的目的來說，關鍵在於它從看起來像 \(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 的線性映射組成的 pair。

這很棒，因為它讓我們一次建立一個 Jacobian 矩陣行，並且評估 \((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

Jacobian-矩陣和矩陣-Jacobian 積

現在你有了 jax.jvp() 和 jax.vjp() 轉換，它們為你提供了函數來一次前推或拉回單個向量，你可以使用 JAX 的 jax.vmap() 轉換來一次前推和拉回整個基底。特別是，你可以使用它來編寫快速的矩陣-Jacobian 和 Jacobian-矩陣積

# 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 的實作

現在我們已經看到了快速的 Jacobian-矩陣和矩陣-Jacobian 積，不難猜測如何編寫 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\) 僅涉及實數輸入和輸出，我們已經知道如何為其編寫 Jacobian-向量積，例如給定一個切向量 \((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 值會揭示梯度（即 Jacobian 或導數）。對於 \(\mathbb{C} \to \mathbb{R}\) 函數，我們可以做同樣的事情：我們仍然可以使用 1.0 作為餘切向量，我們只會得到一個複數結果，總結完整的 Jacobian

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}\) 函數，Jacobian 具有 4 個實值自由度（如上面的 2x2 Jacobian 矩陣中所示），因此我們無法希望在一個複數中表示所有這些自由度。但對於全純函數，我們可以！全純函數恰好是一個 \(\mathbb{C} \to \mathbb{C}\) 函數，其導數可以表示為單個複數。（柯西-黎曼方程式確保上述 2x2 Jacobian 具有複平面中縮放和旋轉矩陣的特殊形式，即單個複數在乘法下的作用。）我們可以透過單次調用具有 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，但我們得到的答案不會代表完整的 Jacobian。相反，它將是函數的 Jacobian，我們只是丟棄輸出的虛部

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. 我們可以透過朝向 grad(f)(x) 的共軛方向邁進，來使用 jax.grad() 來最佳化 \(f : \mathbb{C} \to \mathbb{R}\) 函數，例如複數參數 x 的實值損失函數。

3. 如果我們有一個 \(\mathbb{R} \to \mathbb{R}\) 函數，它恰好在內部使用了一些複數值運算（其中一些運算必須是非全純的，例如卷積中使用的 FFT），那麼 jax.grad() 仍然有效，我們得到的結果與僅使用實數值的實作所給出的結果相同。

在任何情況下，JVP 和 VJP 始終是明確的。如果我們想要計算非全純 \(\mathbb{C} \to \mathbb{C}\) 函數的完整 Jacobian 矩陣，我們可以使用 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\) 的值相乘，這絕非好事。也就是說，我們實際上是在針對大的 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》第三版定義 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\)，因此我們可以認為存在一個函數 \(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 的呼叫來編寫 fixed_point 的 VJP 的方法！此外，在展開 \(A\) 和 \(B\) 後，您可以得出結論，您只需要評估 \(f\) 在 \((a_0, x^*(a_0))\) 處的 VJP 即可。

重點如下

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

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

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

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

def fixed_point_rev(f, res, x_star_bar):
  a, x_star = res
  _, vjp_a = vjp(lambda a: f(a, x_star), a)
  a_bar, = vjp_a(fixed_point(partial(rev_iter, f),
                             (a, x_star, x_star_bar),
                             x_star_bar))
  return a_bar, jnp.zeros_like(x_star)
  
def rev_iter(f, packed, u):
  a, x_star, x_star_bar = packed
  _, vjp_x = vjp(lambda x: f(a, x), x_star)
  return x_star_bar + vjp_x(u)[0]

fixed_point.defvjp(fixed_point_fwd, fixed_point_rev)

print(newton_sqrt(2.))

1.4142135

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

0.35355338
-0.088388346

我們可以透過微分 jnp.sqrt() 來檢查我們的答案，它使用了完全不同的實作方式

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

0.35355338
-0.08838835

這種方法的一個限制是，參數 f 無法封閉任何涉及微分的值。也就是說，您可能會注意到，我們在 fixed_point 的參數列表中保留了明確的參數 a。對於這種使用情況，請考慮使用低階基本運算 lax.custom_root，它允許在具有自訂求根函數的封閉變數中進行微分。

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

使用 jax.custom_jvp 定義前向模式（以及間接地，反向模式）規則

以下是使用 jax.custom_jvp() 的典型基本範例，其中註解使用了 類似 Haskell 的類型簽章

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

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

f.defjvp(f_jvp)

<function __main__.f_jvp(primals, tangents)>

print(f(3.))

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

0.14112
0.14112
-0.9899925

換句話說，我們從一個原始函數 f 開始，它接受類型為 a 的輸入並產生類型為 b 的輸出。我們將其與一個 JVP 規則函數 f_jvp 關聯，該函數接受一對輸入，表示類型為 a 的原始輸入和類型為 T a 的對應切線輸入，並產生一對輸出，表示類型為 b 的原始輸出和類型為 T b 的切線輸出。切線輸出應該是切線輸入的線性函數。

您也可以使用 f.defjvp 作為裝飾器，如下所示

@custom_jvp
def f(x):
  ...

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

即使我們僅定義了 JVP 規則，而沒有定義 VJP 規則，我們也可以在 f 上使用前向和反向模式微分。JAX 將自動轉置我們自訂 JVP 規則中切線值的線性計算，以計算 VJP，其效率就好像我們手動編寫了規則一樣

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

-0.9899925
-0.14112

為了使自動轉置工作，JVP 規則的輸出切線必須是輸入切線的線性函數。否則會引發轉置錯誤。

多個參數的工作方式如下

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

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

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

12.0

defjvps 便利包裝器讓我們可以為每個參數單獨定義 JVP，然後單獨計算結果並將其相加

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

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

print(grad(f)(3.))

-0.9899925

以下是具有多個參數的 defjvps 範例

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

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

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

12.0
12.0
4.0

作為簡寫，使用 defjvps，您可以傳遞 None 值，以指示特定參數的 JVP 為零

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

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

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

12.0
12.0
0.0

呼叫具有關鍵字參數的 jax.custom_jvp() 函數，或編寫具有預設參數的 jax.custom_jvp() 函數定義，都是允許的，只要它們可以根據標準函式庫 inspect.signature 機制檢索的函數簽章，明確地映射到位置參數即可。

當您不執行微分時，函數 f 的呼叫方式就好像它沒有被 jax.custom_jvp() 修飾一樣

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

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

print(f(3.))

called f!
0.14112

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

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

自訂 JVP 規則在微分期間被調用，無論是前向還是反向

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

called f_jvp!
called f!
-0.9899925

print(grad(f)(3.))

called f_jvp!
called f!
-0.9899925

請注意，f_jvp 呼叫 f 以計算原始輸出。在高階微分的上下文中，微分轉換的每次應用都會使用自訂 JVP 規則，當且僅當該規則呼叫原始 f 以計算原始輸出時。（這代表一種基本的權衡，在這種權衡中，我們無法利用規則中 f 評估的中間值，並且讓該規則應用於所有階數的高階微分。）

grad(grad(f))(3.)

called f_jvp!
called f_jvp!
called f!

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

您可以將 Python 控制流程與 jax.custom_jvp() 一起使用

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

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

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

2.0
3.0

使用 jax.custom_vjp 定義自訂的僅限反向模式規則

雖然 jax.custom_jvp() 足以控制前向模式和透過 JAX 的自動轉置進行反向模式微分的行為，但在某些情況下，我們可能希望直接控制 VJP 規則，例如在上面介紹的後兩個範例問題中。我們可以透過 jax.custom_vjp() 來做到這一點。

from jax import custom_vjp

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

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

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

f.defvjp(f_fwd, f_bwd)

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

0.14112
-0.9899925

換句話說，我們再次從一個原始函數 f 開始，它接受類型為 a 的輸入並產生類型為 b 的輸出。我們將其與兩個函數 f_fwd 和 f_bwd 關聯，它們描述了如何分別執行反向模式自動微分的前向和反向傳遞。

函數 f_fwd 描述了前向傳遞，不僅包括原始計算，還包括要儲存哪些值以供反向傳遞使用。它的輸入簽章與原始函數 f 的輸入簽章相同，即它接受類型為 a 的原始輸入。但作為輸出，它產生一個對組，其中第一個元素是原始輸出 b，第二個元素是要儲存以供反向傳遞使用的任何類型為 c 的「殘差」資料。（第二個輸出類似於 PyTorch 的 save_for_backward 機制。）

函數 f_bwd 描述了反向傳遞。它接受兩個輸入，其中第一個是 f_fwd 產生的類型為 c 的殘差資料，第二個是類型為 CT b 的輸出餘切，對應於原始函數的輸出。它產生類型為 CT a 的輸出，表示對應於原始函數輸入的餘切。特別是，f_bwd 的輸出必須是一個序列（例如，一個元組），其長度等於原始函數的參數數量。

因此，多個參數的工作方式如下

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

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

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

f.defvjp(f_fwd, f_bwd)

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

-1.2484405

呼叫具有關鍵字參數的 jax.custom_vjp() 函數，或編寫具有預設參數的 jax.custom_vjp() 函數定義，都是允許的，只要它們可以根據標準函式庫 inspect.signature 機制檢索的函數簽章，明確地映射到位置參數即可。

與 jax.custom_jvp() 一樣，如果未應用微分，則不會調用由 f_fwd 和 f_bwd 組成的自訂 VJP 規則。如果函數被評估，或使用 jax.jit()、jax.vmap() 或其他非微分轉換進行轉換，則只會呼叫 f。

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

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

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

f.defvjp(f_fwd, f_bwd)

print(f(3.))

called f!
0.14112

print(grad(f)(3.))

called f_fwd!
called f!
called f_bwd!
-0.9899925

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

called f_fwd!
called f!
0.14112

print(f_vjp(1.))

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

前向模式自動微分不能用於 jax.custom_vjp() 函數，並且會引發錯誤

from jax import jvp

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

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

如果您想要同時使用前向和反向模式，請改用 jax.custom_jvp()。

我們可以將 jax.custom_vjp() 與 pdb 一起使用，在反向傳遞中插入除錯器追蹤點

import pdb

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

def debug_fwd(x):
  return x, x

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

debug.defvjp(debug_fwd, debug_bwd)

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

jax.grad(foo)(3.)

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

更多功能和細節

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

您應該預期標準 Python 容器（如列表、元組、具名元組和字典）以及它們的巢狀版本都能正常運作。一般來說，任何 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。

對於陣列值參數，您不需要使用 nondiff_argnums，例如，具有整數 dtype 的參數。相反，nondiff_argnums 應僅用於不對應於 JAX 類型（基本上不對應於陣列類型）的參數值，例如 Python 可呼叫物件或字串。如果 JAX 檢測到由 nondiff_argnums 指示的參數包含 JAX Tracer，則會引發錯誤。上面的 clip_gradient 函數是不將 nondiff_argnums 用於整數 dtype 陣列參數的一個很好的範例。

下一步

還有一個充滿其他自動微分技巧和功能的世界。本教學課程未涵蓋但可能值得追求的主題包括

• 高斯-牛頓向量積，線性化一次

• 自訂 VJP 和 JVP

• 定點的高效導數

• 使用隨機 Hessian-向量積估計 Hessian 的跡

• 僅使用反向模式自動微分的前向模式自動微分

• 對於自訂資料類型取導數

• 檢查點機制（用於高效反向模式的二項式檢查點機制，而非模型快照）

• 使用 Jacobian 預先累積優化 VJP