jax.lax.stop_gradient

jax.lax.stop_gradient(x)

停止梯度计算。

在操作上，stop_gradient 是一个恒等函数，也就是说，它返回参数 x 而不改变其值。然而，stop_gradient 会阻止梯度在正向或反向自动微分过程中流动。如果存在多个嵌套的梯度计算，stop_gradient 会停止所有这些计算的梯度。关于其用途的讨论，请参阅 停止梯度。

参数:

x (T) – 数组或数组的Pytree

返回:

输入值保持不变，但在自动微分中将被视为常数。

返回类型:

T

示例

考虑一个返回输入值平方的简单函数

>>> def f1(x):
...   return x ** 2
>>> x = jnp.float32(3.0)
>>> f1(x)
Array(9.0, dtype=float32)
>>> jax.grad(f1)(x)
Array(6.0, dtype=float32)

相同函数，但在 stop_gradient 围绕 x 使用后，在正常评估下结果将相同，但会返回零梯度，因为 x 实际上被视为常数

>>> def f2(x):
...   return jax.lax.stop_gradient(x) ** 2
>>> f2(x)
Array(9.0, dtype=float32)
>>> jax.grad(f2)(x)
Array(0.0, dtype=float32)

这在 JAX 代码库中的许多地方都有使用；例如 jax.nn.softmax() 在内部通过其最大值对输入进行归一化，并且为了效率，这个最大值被封装在 stop_gradient 中。有关 stop_gradient 适用性的更多讨论，请参阅 停止梯度。