jax.vjp#

jax.vjp(fun: Callable[..., T], *primals: Any, has_aux: Literal[False] = False, reduce_axes: Sequence[AxisName] = ()) → tuple[T, Callable][源代码]#
jax.vjp(fun: Callable[..., tuple[T, U]], *primals: Any, has_aux: Literal[True], reduce_axes: Sequence[AxisName] = ()) → tuple[T, Callable, U]

计算 fun 的（反向模式）向量-雅可比积。

grad() 被实现为 vjp() 的一个特例。

参数：

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

  • primals – 一个原始值的序列，应该在该序列处评估 fun 的雅可比矩阵。 primals 的数量应等于 fun 的位置参数的数量。每个原始值都应是一个数组、标量或它们的 pytree（标准 Python 容器）。

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

返回值：

如果 has_aux 为 False，则返回一个 (primals_out, vjpfun) 对，其中 primals_out 是 fun(*primals)。如果 has_aux 为 True，则返回一个 (primals_out, vjpfun, aux) 元组，其中 aux 是 fun 返回的辅助数据。

vjpfun 是一个函数，它从一个与 primals_out 具有相同形状的余切向量映射到一个与 primals 具有相同数量和形状的余切向量的元组，表示在 primals 处评估的 fun 的向量-雅可比积。

>>> import jax
>>>
>>> def f(x, y):
...   return jax.numpy.sin(x), jax.numpy.cos(y)
...
>>> primals, f_vjp = jax.vjp(f, 0.5, 1.0)
>>> xbar, ybar = f_vjp((-0.7, 0.3))
>>> print(xbar)
-0.61430776
>>> print(ybar)
-0.2524413