torch.func.grad_and_value

torch.func.grad_and_value(func, argnums=0, has_aux=False)

返回一个函数，用于计算梯度和原始值或前向计算的元组。

参数

• func (Callable) – 一个 Python 函数，它接受一个或多个参数。必须返回一个单元素张量。如果指定 has_aux 等于 True，则函数可以返回一个包含单元素张量和其他辅助对象的元组：(output, aux)。

• argnums (int 或 Tuple[int]) – 指定要计算关于哪些参数的梯度。argnums 可以是单个整数或整数元组。默认值：0。

• has_aux (bool) – 标志，指示 func 返回一个张量和其他辅助对象：(output, aux)。默认值：False。

返回

用于计算关于其输入的梯度元组和前向计算的函数。默认情况下，函数的输出是关于第一个参数的梯度张量和一个原始计算的元组。如果指定 has_aux 等于 True，则返回梯度元组和包含输出辅助对象的前向计算元组。如果 argnums 是整数元组，则返回关于每个 argnums 值的输出梯度元组和前向计算。

返回类型

Callable

有关示例，请参见 grad()