See #347 for example code.

Interestingly, Autograd still has this issue:

import autograd.numpy as np
from autograd import grad

def f(x):
  x = np.sqrt(np.sum(x**2, axis=1))
  return np.sum(np.where(x > 0.5, x, np.ones_like(x)))

x = np.array([[1, 2], [3, 4], [0, 0]], dtype=np.float64)
print grad(f)(x)

# [[0.4472136  0.89442719]
#  [0.6        0.8       ]
#  [       nan        nan]]

It comes down to the formula for the jvp of sqrt(x) being problematic at x=0 when the gradient has coefficients that are also zero, since g * x ** -0.5 is 0 * inf there. In this case, the zero coefficients in the gradient arise from the lax.select/np.where call.

I think a zero in the gradient should win out over an inf in the Jacobian, and this PR implements that policy. I had trouble implementing it on top of the standard lax.mul primitive because of linearity constraints in the JVP rule, so instead I added a new lax._safe_mul that implements the 0*inf = 0 policy in its translation rule.