jax.value_and_grad#

jax.value_and_grad(fun, argnums=0, has_aux=False, holomorphic=False, allow_int=False, reduce_axes=())[源代码]#

创建一个函数，该函数既评估 fun 又评估 fun 的梯度。

参数:

  • fun (可调用对象) – 需要进行微分的函数。其由 argnums 指定的位置上的参数应为数组、标量或标准 Python 容器。它应返回一个标量（包括形状为 () 的数组，但不包括形状为 (1,) 等的数组）。

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

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

  • holomorphic (bool) – 可选，布尔值。指示 fun 是否保证是全纯的。如果为 True，则输入和输出必须是复数。默认为 False。

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

  • reduce_axes (Sequence[AxisName])

返回:

一个与 fun 具有相同参数的函数，该函数同时计算 fun 和 fun 的梯度，并将它们作为一对（一个包含两个元素的元组）返回。如果 argnums 是一个整数，则梯度具有与该整数指示的位置参数相同的形状和类型。如果 argnums 是一个整数序列，则梯度是一个元组，其值的形状和类型与相应的参数相同。如果 has_aux 为 True，则返回一个 ((值, 辅助数据), 梯度) 的元组。

返回类型:

Callable[…, tuple[Any, Any]]