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]),\]

表示在 \(x\) 处的 \(h\) 的 \(K\) 阶泰勒近似，jet() 返回在 \(x\) 处的 \(f\) 的 \(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)[源代码][源代码]#

泰勒模式高阶自动微分。

参数:

  • 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
目录