NaN when computing gradient of squared norm evaluated at 0. about jax HOT 2 CLOSED

Hi, thanks for the report! This is working as intended: the gradient of the norm at 0.0 is not well-defined, so NaN is the most reasonable output.

Mathematically, you have

$norm(x) = \left[\sum_i x_i^2\right]^{1/2}$

The partial derivative with respect to each input is

$d/dx_i norm(x) = x_i \left[\sum_i x_i^2\right]^{-1/2} = x_i/norm(x)$

And if you evaluate this at $x=0$ you get $0/0$, which is undefined.

from jax.

That is true. For squared norm however, it would be $2x_i$. I can instead use jnp.sum(x**2) to avoid the issue. Thanks!

from jax.