Skip to content

← JAX Autodiff step 11 of 25

Medium Primitives

Custom JVP: Clip with Pass-through Gradient

Jump to the task

Clipping keeps values inside an interval. Outside that interval, changing the input slightly does not change the clipped output, so the ordinary derivative with respect to that input is zero.

This exercise deliberately substitutes a different derivative: treat the clipping operation as the identity when propagating derivatives with respect to x. This is a straight-through estimator, or STE. It is a surrogate rule, not the mathematical derivative of clipping.

Forward values and tangents

jax.custom_jvp lets you supply a Jacobian-vector product rule. Its arguments are tuples of primal values and tangent values, one entry per function argument. Return both the primal output and its tangent:

@jax.custom_jvp
def clip_passthrough(x, lo, hi):
    return jnp.clip(x, lo, hi)

@clip_passthrough.defjvp
def clip_jvp(primals, tangents):
    x, lo, hi = primals
    x_dot, lo_dot, hi_dot = tangents
    return jnp.clip(x, lo, hi), x_dot

The forward output still clips. The tangent is exactly x_dot, including for inputs below, above, or equal to the bounds. The rule is linear in its tangent arguments, as required for JAX to transpose it into a reverse-mode rule.

Here lo_dot and hi_dot are intentionally ignored. Scalars are not inherently constants to autodiff: if someone differentiates with respect to either bound, this custom rule gives zero. That is part of this exercise's chosen surrogate, not a property of ordinary clipping.

A concrete comparison

For x = [-2, 0.5, 3] and bounds [0, 1], the output is [0, 0.5, 1]. With an input tangent [2, -3, 4], this custom JVP returns tangent [2, -3, 4]. Ordinary clipping would return [0, -3, 0] for those inputs, which are away from the boundary points.

For the scalar objective sum(clip_passthrough(x, 0, 1)), jax.grad returns [1, 1, 1]. More generally, an upstream cotangent passes through unchanged; gradients are not always ones. An outer loss can change them.

STEs are used in some quantization and discrete-model training methods, but this rule does not guarantee useful updates or convergence. It does not change the forward objective into a smoother function. Finite differences of the clipped output will disagree with the surrogate outside the interval, by design.

Problem

Implement clip_passthrough(x, lo, hi) with @jax.custom_jvp:

  1. Return jnp.clip(x, lo, hi) as the primal output.
  2. Register a JVP rule with defjvp.
  3. Return (primal_output, x_dot) from that rule, ignoring bound tangents.

x is a finite float32 JAX vector of shape (n,), with n >= 1. lo and hi are finite real scalars with lo <= hi. Return a JAX vector of the same shape and leave x unchanged. Equal bounds are allowed: the output is constant, while the chosen tangent still passes through.

Output tests check clipping values only, so a plain clip can pass them without the requested derivative rule. Check jax.jvp and a weighted scalar objective with jax.grad locally to verify the STE behavior.