Raise when differentiating functions at a non-differentiable point?

[cvjjm]

I know that for functions that are mathematically not everywhere differentiable (for example np.abs()) JAX tries to return a "reasonable" gradient anyway, i.e. jax.grad(np.abs)(0.) returns 1. (and I think this is useful).

But now I have a (jit compiled) function (much more complicated and thus expensive to evaluate and automatically differentiate than abs) that I also know does not have a well defined gradient in certain points. I would like to raise an exception and handle such cases in the calling code rather than having the function return a useless (because it is unstable) automatic derivative after a lot of expensive computation.

In autograd I could detect when my code is traced for differentiation with an ArrayBox and raise an exception.

Is there a way to do the same in JAX? Can I somehow put code into my function that doesn't get in the way of (and actually survives) jit compilation, allows the function to be evaluated everywhere normally, but makes it raise an exception when its derivative is calculated at certain points? Thanks!

[JiaYaobo]

jax.pure_callback may work, but it can't show raise statement, instead you can print it out in callback, like

import jax
import jax.numpy as jnp

@jax.custom_jvp
def absolute(x):
    return jnp.where(x >= 0, x, -x)

def assert_at0(x):
    if x == 0.:
        print("x is not derivable at 0")
        raise AssertionError("x is not derivable at 0")
    return x

@absolute.defjvp
def absolute_jvp(primals, tangents):
    x, = primals
    shape_dtype = jax.ShapeDtypeStruct(x.shape, x.dtype)
    x = jax.pure_callback(assert_at0, shape_dtype, x)
    x_dot, = tangents
    return absolute(x), x_dot * jnp.sign(x)

jax.grad(jax.jit(absolute))(jnp.array(0.))

[cvjjm]

I can raise from a callback called with jax.debug.callback(), but this method would require me to implement a custom jvp for my function. Since JAX can handle the auto differentiation of my function just fine I would rather not want to do that. Is there a way to avoid marking my function custom_jvp or jax.custom_vjp? Or alternatively a way to "resume" the standard tracing and differentiation logic inside a custom jvp or vjp?

[jakevdp]

The best way to do these kinds of runtime checks is with the experimental checkify transformation. For example:

import jax
import jax.numpy as jnp
from jax.experimental import checkify

@checkify.checkify
def grad_abs(x):
  checkify.check(x!=0, "x must be nonzero", x=x)
  return jax.grad(jnp.abs)(x)

err, out = jax.jit(grad_abs)(0.)
err.throw()

---------------------------------------------------------------------------
FailedCheckError                          Traceback (most recent call last)
FailedCheckError: x must be nonzero (check failed at <ipython-input-2-239f912d24b1>:7 (grad_abs))

The above exception was the direct cause of the following exception:
...
JaxRuntimeError: x must be nonzero (check failed at <ipython-input-2-239f912d24b1>:7 (grad_abs))

The API is still experimental and subject to change, but this is the main mechanism available for this kind of check.

[JiaYaobo]

chekify is cool! never use it before!

[cvjjm]

This would work (I would have to take my function instead of jnp.abs of course) but only as long as everybody calls the checkified grad function instead of applying jax.grad directly to my function or a function calling my function.

[cvjjm]

The workaround I settled on in the end is to make an "identity" function that is not everywhere differentiable, i.e., something like

    @jax.custom_jvp
    @jax.jit
    def not_everywhere_differentiable(y):
        return y

    @not_everywhere_differentiable.defjvp
    @jax.jit
    def not_everywhere_differentiable_jvp(primals, tangents):
        from jax import lax

        def raise_not_differentiable(primals):
            raise NotDifferentiableError(primals)

        def call_raise_not_differentiable(primals):
            jax.debug.callback(raise_not_differentiable, primals)

        y = primals[0]
        lax.cond(
            jax.numpy.isclose(y, 0.),
            call_raise_not_differentiable,
            lambda _: None,
            primals,
        )

        return y, tangents[0]

and use that somewhere in my custom function to transform a variable y = not_everywhere_differentiable(y).

In this way the code above ends up in the automatically generated gradient code of my function and I don not need to provide a manual vjp for my complicated function.

Still, it would be nice if JAX offered a way to "attach" some custom code to the gradient computation of a function without having to define a custom vjp.