Jax.lax.stop_gradient()

As explained in Deep Learning with JAX

A low-level JAX operation that passes values through unchanged while blocking gradient propagation through them.

jax.lax.stop_gradient() passes a value through unchanged while blocking gradient flow through that path. It is useful when only part of a computation should affect autodiff, such as excluding one branch from an update. section 4.2.3

Applied to an input, it behaves like an identity in the forward pass, but reverse-mode and forward-mode differentiation treat it as non-differentiable. The result is zero gradient for that path. section 4.2.3

Blocking one branch

For f(x, y) = x**2 + stop_gradient(y**2), differentiating at (1.0, 1.0) gives gradients (2.0, 0.0). The x branch still contributes; the y branch is ignored by autodiff. section 4.2.3

Use it when gradient control matters, especially in these settings:

• Exclude variables from one loss while optimizing another. section 4.2.3
• Prevent a branch in a computational graph from influencing updates. section 4.2.3
• Handle cases common in reinforcement learning. section 4.2.3

Figure 4.3 Obtaining function for calculating gradients

The computational-graph figure makes the branch-selective gradient behavior easy to see. section 4.2.3

How Deep Learning with JAX covers it

How deep it goes

Defines identity behavior plus autodiff blocking, and shows the resulting zero-gradient effect in a worked example.

Examples

Demonstrates a two-branch x**2 + y**2 example with one branch wrapped in stop_gradient().

What it assumes

Assumes basic gradient calculation and introduces the blocking behavior from scratch.

What you can do after

Apply it to keep selected values from affecting gradient updates.

How it connects

Relates to autodiff, computational graphs, and PyTorch detach().

liveBookmanning.com

Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

• jax.gradcontrasts with · jax.grad is a JAX transformation that creates a function computing derivatives, typically gradients of a scalar-output function with respect to specified inputs.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.