Gradient of arctan2 fails but unclear why

[Peter-Vincent]

I have a large network that, at some point, uses the jnp.arctan2(y,x) function to turn coordinates into angles. I need to backprop through this expression, however, I invariably see this error

FloatingPointError: invalid value (nan) encountered in jit(div)

at the line containing the call to arctan2, with config.update("jax_debug_nans",True)

The function performs fine in the forward pass and does not produce any nan values.

Looking at the gradient of arctan2 this would make sense if both x and y = 0 (or at least very very small) (the partial derivatives are -y/(y**2 + x**2) and x/(y**2 + x**2), but if I run the following minimal example

def f(x,y):
    return jnp.sum(jnp.arctan2(y,x))

if __name__ == "__main__":
    grad_f = value_and_grad(f,argnums=(0,1))
    x = jnp.array([1.0,-1.0,0.5,-0.5])
    y = jnp.array([-1.0,0.5,-0.5,1.0])
    print(grad_f(y,x))
    # Now set the x values to 0
    x = jnp.array([0.0,0.0,0.0,0.0])
    print(grad_f(y,x))
    y = jnp.array([0.0,0.5,0.0,1.0])

the gradients are always evaluated well. It's a bit tricky to know how to start resolving this problem, since clearly my diagnosis is wrong. What are the next steps to trying to understand how this issue is arising?

Going through the full code in debug mode and breaking when we reach that line, it doesn't seem like any of the values are 0, but I'm still struggling to peer into the Jax internals.

Adding some stuff

Printing the jaxpr gives me the following

import jax
import jax.numpy as jnp

def a(y,x):
    return jnp.arctan2(y,x)

x = jnp.float32(0.7)
y = jnp.float32(0.3)

print(jax.make_jaxpr(a)(y,x))
# { lambda ; a:f32[] b:f32[]. let c:f32[] = atan2 a b in (c,) }
print(jax.make_jaxpr(jax.grad(a))(y,x))
#{ lambda ; a:f32[] b:f32[]. let
#    _:f32[] = atan2 a b
#    c:f32[] = integer_pow[y=2] a
#    d:f32[] = integer_pow[y=2] b
#    e:f32[] = add c d
#    f:f32[] = div b e
#    g:f32[] = mul 1.0 f
#  in (g,) }

so it all seems good until that div step where I expect the error is coming from, but I don't quite understand why my minimal example should not fail?

[jakevdp]

Thanks for the question!

Operationally, the issue is that the gradient of atan2(x, y) contains the term x / (x ** 2 + y ** 2), so if x and y have zeros in the same location, the gradient will be NaN.

More conceptually, atan2 is not continuous at the point $(0, 0)$, so its derivative at that point is ill-defined.

[Peter-Vincent]

So it was what I figured, and the reason my minimal example was passing is because I accidentally deleted the final calcualtion of the gradient after setting y, silly me. Thanks for answering and confirming!