Clipping changing the gradients

[SNMS95]

import jax.numpy as jnp
import jax

def loss(x, with_clip=True):
    if with_clip:
        x = jnp.clip(x, 1e-3, 1.0)
    return jnp.sum(x**2)

x = jnp.array([[1.0, 1.0, 1.0],
               [0.01, 0.01, 1.0]])

# jnp.array([[0.9, 0.9, 0.9],
#                [0.01, 0.01, 0.9]])
print("Loss:", loss(x))
jax.grad(loss)(x)

In this code, if the clip operation does not do anything to the inputs. However, if we check the gradient with and without the clipping operation, we get two different values. Is this because max and min functions, which are used behind the scenes, get a gradient of 0.5 for the edge cases? This goes away, if we choose the input to not include the upper and lower limits!

[jakevdp]

Hi - thanks for the question! I think this is working more-or-less as expected. clip is a function with a discontinuous gradient, and you're asking for the gradient at the precise point of the discontinuity. There are four reasonable choices for the gradient of clip(x, 0, 1) evaluated at x = 1.0:

1. taking a symmetric limit would give you NaN
2. taking the limit from the left would give you 1.0 (this would result in both your gradients being equal)
3. taking the limit from the right would give you 0.0 (this is what is done by jax.lax.clamp for the same case)
4. to make the gradient equivalent under argument permutation, we can return 0.5 (this is what is done by jnp.clip, via the autodiff rule of jax.lax.min and jax.lax.max)

From a strict mathematical perspective, (1) might be the best, but introduing NaNs into autodiff code can be problematic. So we choose one of the other three options, each of which are reasonable choices. What do you think?

[jakevdp]

Side note: I was wondering why jnp.clip doesn't just call into lax.clamp; it turns out that's because lax.clamp doesn't support complex numbers, and you recover correct semantics for both real and complex clip by implementing it in terms of jnp.min and jnp.max (see #225).

[SNMS95]

I think this is clear but it was a quite a surprise to actually see it. (Also did you mean that option 1 is the best from a math perspective? since that's the only one that would give nans)
Would it be better to document this in the documentation of these functions itself?
Or in the AD tutorial that this is convention being followed?

[jakevdp]

Yes, I meant option 1 (edited to update).

I don't know of any documentation of specific edge-case behavior for autodiff of particular functions: I imagine trying to document that exhaustively for all APIs would be a never-ending rabbit-hole. And there's not really any globally-true convention: sometimes choices are motivated by mathematics or symmetry (e.g. in the case of max) sometimes they're an accident of implementation (e.g. in the case of clip) and sometimes they're motivated by historical precedent / decisions used in the literature (e.g. as is often the case for jax.nn APIs).