nan when using np.where and grad

Here is the situation:

f = lambda x: np.where(x<1., x, np.sqrt(x))
df = grad(f)
df(1e-9)  # evaluates to 1, as expected
df(0.0)   # evaluates to nan, presumably because the second case of the where is nan in that case
         # but it is the first case that should give the result, not the second!

Not sure if this is expected behavior, but it is surprising. I’m going to work on a work-around, but thought I should document it. Thanks!

[mrtupek]

Interesting… I see the dilemma now.

I think for my usage, I have to conclude that the wrong derivative of sqrt(x)^2 at x=0 is probably the lesser of two evils. Especially because my actual use case is to construct a ‘smooth’ sqrt(x dot x) which is quadratic around x=0, so I’ll never actually evaluate it there. Thanks a lot for the help and the links.

[mattjj]

There is something of a gotcha, but it’s not obvious! It’s related to the one mentioned in that comment on #1052 that Jake linked.

print(grad(lambda x: safe_sqrt(x)**2)(0.))  # prints 0.0

I think that should be nan because the function isn’t differentiable over the reals at 0. It is differentiable over the nonnegative reals, but there the derivative is 1.0, not 0.0 as computed. (Related: enforcing differentiation conventions using custom_jvp.)