#802·autograd

Gradient of where() is sometimes erroneously NaN when the inactive case is NaN

Author: jkunimuneCreated Aug 30, 2026Updated Aug 30, 2026

Functions defined with where() often have gradients of NaN when they should be defined. It seems to happen when the case that's not chosen is NaN. For example, consider the following simple piecewise function. Even though the second case is only defined for negative x, the overall function should be defined and differentiable everywhere except maybe at x=0. But instead it evaluates to NaN when x>0.

python
from autograd import jacobian
from autograd.numpy import where, sqrt

def f(x):
	return where(x >= 0, x, sqrt(-x))

f_prime = jacobian(f)

print(f(1.))  # expected and actual: 1.0
print(f_prime(1.))  # expected: 1.0; actual: nan

Note that it doesn't happen all the time when the second case is NaN. For example, if you replace sqrt with log above, the problem goes away and you get the expected answer. Or if you change where(x >= 0, x, sqrt(-x)) to where(x <= 0, x, sqrt(x)), the problem goes away. I don't know why sqrt(-x) is uniquely troubled.