NaN values while using jax.numpy.sqrt

[aonurdasdemir]

I am implementing a sobel filter inside my optimization package and I need to use vjp on that function. However, since sobel filter has a term G = jnp.sqrt(jnp.square(G_x) + jnp.square(G_y) where G_x and G_y are arrays and some elements of them are zero. When I use vjp on this function, it results in a array that has Nan values. From what I see on previous questions, the problem is using sqrt with zero values. Are there any possible solutions or workarounds to solve this problem?

Here is a minimized code which results in nan in vjp for functions that has sqrt and input array has zero in some elements:

def test_func(a):
    result = jnp.sqrt(jnp.square(a))
    return result

a = jnp.array([[0.0, 1.0], [1.0, 1.0]])

primals, f_vjp = vjp(test_func, a)
vjp_result = f_vjp(1.0)

[jakevdp]

The best way to deal with this is probably via custom JVP/VJP rules. This situation is very similar to the Numerical Stability section here: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#numerical-stability

If you define an appropriate custom JVP/VJP for your function, you can control the behavior at values that are problematic for the raw floating point implementation.

[mattjj]

Maybe an even closer fit would be the differentiation convention section, since IIUC sqrt is differentiable over [0, inf) but not over R.