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, factorial_scaled=True, **_)[源代码]#

泰勒模式高阶自动微分。

参数:

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

  • primals – 计算 fun 泰勒近似时所处的原值（primal values）。应为参数的元组或列表，其长度应等于 fun 的位置参数数量。

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

  • factorial_scaled – 若为 True，则输入和输出序列中的每一项都按其阶数的阶乘进行缩放，从而使输入和输出序列成为泰勒级数。这是默认行为，使得输入和输出序列中的第 n 阶项即为函数的第 n 阶导数。若为 False，则输入和输出序列为未经过阶乘缩放的泰勒系数（即泰勒级数中每一项的常数系数）。

返回:

一个 (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.74414825 0.74414825

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