jax.linearize

内容

jax.linearize#

jax.linearize(fun: Callable, *primals, has_aux: Literal[False] = False) → tuple[Any, Callable][source]#
jax.linearize(fun: Callable, *primals, has_aux: Literal[True]) → tuple[Any, Callable, Any]

使用 jvp() 和部分求值，产生函数 fun 的线性近似。

参数:

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

  • primals – 应该计算 fun 的雅可比矩阵的原值。应该是一个数组元组、标量或其标准 Python 容器。元组的长度等于 fun 的位置参数数量。

  • has_aux – 可选，布尔值。指示 fun 是否返回一对，其中第一个元素被认为是线性化数学函数的输出，第二个元素是辅助数据。默认值为 False。

返回值:

如果 has_aux 为 False，则返回一个对，其中第一个元素是 f(*primals) 的值，第二个元素是一个函数，该函数评估在 primals 处评估的 fun 的（前向模式）雅可比-向量乘积，无需重新执行线性化工作。如果 has_aux 为 True，则返回一个 (primals_out, lin_fn, aux) 元组，其中 aux 是 fun 返回的辅助数据。

在计算的值方面，linearize() 的行为类似于一个柯里化的 jvp()，其中这两个代码块计算相同的值

y, out_tangent = jax.jvp(f, (x,), (in_tangent,))

y, f_jvp = jax.linearize(f, x)
out_tangent = f_jvp(in_tangent)

但是，区别在于 linearize() 使用部分求值，因此函数 f 在调用 f_jvp 时不会被重新线性化。一般来说，这意味着内存使用量会随着计算量的增长而增长，这与反向模式非常相似。（实际上，linearize() 具有与 vjp() 相似的签名！）

如果您想多次应用 f_jvp，例如，为相同线性化点处的许多不同输入切向量评估前推，则此函数非常有用。此外，如果所有输入切向量都已知，则可以使用 vmap() 对其进行向量化，效率更高，如

pushfwd = partial(jvp, f, (x,))
y, out_tangents = vmap(pushfwd, out_axes=(None, 0))((in_tangents,))

通过使用 vmap() 和 jvp() 结合使用，我们可以避免 linearize() 和 vjp() 所产生的与计算深度成正比的存储线性化内存成本。

这是一个更完整的使用 linearize() 的示例

>>> import jax
>>> import jax.numpy as jnp
>>>
>>> def f(x): return 3. * jnp.sin(x) + jnp.cos(x / 2.)
...
>>> jax.jvp(f, (2.,), (3.,))
(Array(3.26819, dtype=float32, weak_type=True), Array(-5.00753, dtype=float32, weak_type=True))
>>> y, f_jvp = jax.linearize(f, 2.)
>>> print(y)
3.2681944
>>> print(f_jvp(3.))
-5.007528
>>> print(f_jvp(4.))
-6.676704
目录