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)

在 x 周围使用 stop_gradient 的相同函数在正常求值下是等效的，但会返回零梯度，因为 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 的适用性的更多讨论，请参考 停止梯度。