Repository: google/jax
Issue: Issue taking gradients through np.where when one of branches is nan.

If you take gradients through np.where and one of the branches is nan, the gradient will be nan even if the branch is not taken. 

Here's a repro:


```
def test_fail(x):
  x = np.sqrt(np.sum(x ** 2, axis=1))
  ones = np.ones_like(x)
  x = np.where(x > 0.5, x, ones)
  return np.sum(x)

def test_succeed(x):
  x = np.sum(x ** 2, axis=1)
  ones = np.ones_like(x)
  x = np.sqrt(np.where(x > 0.5, x, ones))
  return np.sum(x)

x = np.array([[1, 2], [3, 4], [0, 0]], dtype=np.float64)

print('Test np.sqrt then np.where = {}'.format(test_fail(x)))
print('Gradient:')
print(grad(test_fail)(x))

print('Test np.where then np.sqrt = {}'.format(test_succeed(x)))
print('Gradient:')
print(grad(test_succeed)(x))
```


Changed files:
jax/lib/xla_bridge.py, jax/lax.py, tests/lax_numpy_test.py

diff --git a/jax/lax.py b/jax/lax.py
index 2c6c62b13..b75b521c9 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -493,6 +493,9 @@ def stop_gradient(x):
   return stop_gradient_p.bind(x)
 
 
+def _safe_mul(x, y): return safe_mul_p.bind(x, y)
+
+
 def psum(x, axis_name):
   return psum_p.bind(x, axis_name=axis_name)
 
@@ -957,11 +960,18 @@ _maybe_real = lambda x: real(x) if _iscomplex(x) else x
 
 # TODO handle broadcasting
 pow_p = standard_binop([_float | _complex, _float | _complex], 'pow')
-ad.defjvp(pow_p,
-          lambda g, x, y: mul(_brcast(g, y), mul(y, pow(x, select(
-              eq(y, _zeros(y)), _ones(y), sub(y, _ones(y)))))),
-          lambda g, x, y: mul(_brcast(g, x),
-                              mul(log(_replace_zero(x)), pow(x, y))))
+
+def pow_jvp_lhs(g, x, y):
+  # we call _safe_mul here so that we get the behavior 0*inf = 0, since when a
+  # coefficient in `g` is zero we want to keep it at zero, not produce a nan.
+  # see https://github.com/google/jax/pull/383
+  jac = mul(y, pow(x, select(eq(y, _zeros(y)), _ones(y), sub(y, _ones(y)))))
+  return _safe_mul(_brcast(g, y), jac)
+
+def pow_jvp_rhs(g, x, y):
+  return mul(_brcast(g, x), mul(log(_replace_zero(x)), pow(x, y)))
+
+ad.defjvp(pow_p, pow_jvp_lhs, pow_jvp_rhs)
 _replace_zero = lambda x: select(eq(x, _const(x, 0)), _ones(x), x)
 
 not_p = standard_unop(_int | _bool, 'not')
@@ -992,6 +1002,20 @@ mul_p = standard_binop([_num, _num], 'mul')
 ad.defbilinear_broadcasting(_brcast, mul_p, mul, mul)  # TODO
 
 
+def _safe_mul_translation_rule(c, x, y):
+  dtype = c.GetShape(x).numpy_dtype()
+  zero = c.Constant(onp.array(0, dtype=dtype))
+  out_shape = tuple(onp.maximum(c.GetShape(x).dimensions(),
+                                c.GetShape(y).dimensions()))
+  return c.Select(c.Or(c.Eq(x, zero), c.Eq(y, zero)),
+                  c.Broadcast(zero, out_shape),
+                  c.Mul(x, y))
+
+safe_mul_p = standard_binop([_num, _num], 'safe_mul',
+                            translation_rule=_safe_mul_translation_rule)
+ad.defbilinear_broadcasting(_brcast, safe_mul_p, _safe_mul, _safe_mul)
+
+
 def _div_transpose_rule(cotangent, x, y):
   assert x is None and y is not None
   res = ad_util.zero if cotangent is ad_util.zero else div(cotangent, y)
diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
index f35748efb..b19c04a11 100644
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -348,7 +348,7 @@ class _JaxComputationBuilderBase(object):
 
   def ConstantLike(self, example_value, value, canonicalize_types=True):
     example_value = onp.asarray(example_value)
-    return self.Constant(onp.array(value).astype(example_value.dtype))
+    return self.Constant(onp.array(value, dtype=example_value.dtype))
 
   def Constant(self, py_val, canonicalize_types=True):
     """Translate constant `py_val` to a constant for this ComputationBuilder.
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 863a025f1..869e2ac92 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1223,6 +1223,18 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     jax_numpy_result = ((x + x.T) / 2).dtype
     self.assertEqual(orig_numpy_result, jax_numpy_result)
 
+  def testIssue347(self):
+    # https://github.com/google/jax/issues/347
+    def test_fail(x):
+      x = lnp.sqrt(lnp.sum(x ** 2, axis=1))
+      ones = lnp.ones_like(x)
+      x = lnp.where(x > 0.5, x, ones)
+      return lnp.sum(x)
+
+    x = lnp.array([[1, 2], [3, 4], [0, 0]], dtype=lnp.float64)
+    result = api.grad(test_fail)(x)
+    assert not onp.any(onp.isnan(result))
+
 
 if __name__ == "__main__":
   absltest.main()