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))
```



Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- 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 @@ def _conj_transpose_rule(t, x, input_dtype):
 
 # 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 @@ def _add_transpose(t, x, y):
 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
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -348,7 +348,7 @@ def NumpyArrayConstant(self, value, canonicalize_types=True):
 
   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.


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1223,6 +1223,18 @@ def testSymmetrizeDtypePromotion(self):
     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()