jax.jacrev

jax.jacrev(fun, argnums=0, has_aux=False, holomorphic=False, allow_int=False)

函数fun的雅可比矩阵，使用反向模式自动微分逐行计算得到。

参数:

• fun (Callable) – 要计算其雅可比矩阵的函数。

• argnums (int | Sequence[int]) – 可选，整数或整数序列。指定要对其进行微分的位置参数（默认为0）。

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

• holomorphic (bool) – 可选，布尔值。指示fun是否被认为是全纯的。默认为False。

• allow_int (bool) – 可选，布尔值。是否允许对整数值输入进行微分。整数输入的梯度将具有一个平凡的向量空间数据类型（float0）。默认为False。

返回:

一个与fun具有相同参数的函数，它使用反向模式自动微分计算fun的雅可比矩阵。如果has_aux为True，则返回一个对(jacobian, auxiliary_data)。

返回类型:

可调用对象

>>> import jax
>>> import jax.numpy as jnp
>>>
>>> def f(x):
...   return jnp.asarray(
...     [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])])
...
>>> print(jax.jacrev(f)(jnp.array([1., 2., 3.])))
[[ 1.       0.       0.     ]
 [ 0.       0.       5.     ]
 [ 0.      16.      -2.     ]
 [ 1.6209   0.       0.84147]]