jnp.where and NaN in gradients

[PgLoLo]

Consider the following code snippet:

from jax import value_and_grad, numpy as jnp


def bad_func(x):
    return (x**2 + jnp.inf)**2

def good_func(x):
    return jnp.where(x > 0, x, bad_func(x))

print(value_and_grad(good_func)(.1))

Here we created a function bad_func, gradient for which is undefined, and a function good_func that on R+ is just x. Unfortunately, autograd stragles to calculate the gradient of the given function in the point 0.. The output of the presented code would be (DeviceArray(0.1, dtype=float32), DeviceArray(nan, dtype=float32)).

That's very unfortunate. Any suggestions on how to workaround this issue?

[jakevdp]

There is some information on this topic in the Frequently Asked Questions: https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where