jax.experimental.jet 模块

内容

jax.experimental.jet 模块#

Jet 是一个用于更高阶自动微分的实验模块，它不依赖于重复的一阶自动微分。

如何？通过截断泰勒多项式的传播。考虑一个函数 \(f = g \circ h\)，某个点 \(x\) 和某个偏移 \(v\)。一阶自动微分（例如 jax.jvp()）从对 \((h(x), \partial h(x)[v])\) 计算对 \((f(x), \partial f(x)[v])\)。

jet() 实现更高阶的模拟：给定元组

\[(h_0, ... h_K) := (h(x), \partial h(x)[v], \partial^2 h(x)[v, v], ..., \partial^K h(x)[v,...,v]),\]

它代表着 \(h\) 在 \(x\) 处的 \(K\) 阶泰勒近似，jet() 返回 \(f\) 在 \(x\) 处的 \(K\) 阶泰勒近似，

\[(f_0, ..., f_K) := (f(x), \partial f(x)[v], \partial^2 f(x)[v, v], ..., \partial^K f(x)[v,...,v]).\]

更具体地说，jet() 计算

\[f_0, (f_1, . . . , f_K) = \texttt{jet} (f, h_0, (h_1, . . . , h_K))\]

因此可用于对 \(f\) 进行高阶自动微分。详细内容请参阅 这些笔记。

注意

通过贡献 未完成的原始规则 来帮助改进 jet()。

API#

jax.experimental.jet.jet(fun, primals, series)[source]#

泰勒模式高阶自动微分。

参数:

  • fun – 要微分的函数。其参数应为数组、标量或数组或标量的标准 Python 容器。它应返回数组、标量或数组或标量的标准 Python 容器。

  • primals – fun 泰勒近似应在其中求值的原始值。应为参数元组或列表，其长度应等于 fun 的位置参数数量。

  • series – 高阶泰勒级数系数。primals 和 series 共同构成了截断的泰勒多项式。应为元组或列表的元组或列表，其长度决定了截断的泰勒多项式的次数。

返回值:

一个 (primals_out, series_out) 对，其中 primals_out 为 fun(*primals)，并且 primals_out 和 series_out 共同构成了 \(f(h(\cdot))\) 的截断泰勒多项式。primals_out 值与 primals 具有相同的 Python 树结构，而 series_out 值与 series 具有相同的 Python 树结构。

例如

>>> import jax
>>> import jax.numpy as np

考虑函数 \(h(z) = z^3\)，\(x = 0.5\)，以及前几个泰勒系数 \(h_0=x^3\)，\(h_1=3x^2\) 和 \(h_2=6x\)。令 \(f(y) = \sin(y)\)。

>>> h0, h1, h2 = 0.5**3., 3.*0.5**2., 6.*0.5
>>> f, df, ddf = np.sin, np.cos, lambda *args: -np.sin(*args)

jet() 根据 Faà di Bruno 公式返回 \(f(h(z)) = \sin(z^3)\) 的泰勒系数

>>> f0, (f1, f2) =  jet(f, (h0,), ((h1, h2),))
>>> print(f0,  f(h0))
0.12467473 0.12467473

>>> print(f1, df(h0) * h1)
0.7441479 0.74414825

>>> print(f2, ddf(h0) * h1 ** 2 + df(h0) * h2)
2.9064622 2.9064634
目录