Repository: google/jax
Issue: Gradient of `tanh` sometimes causes invalid values

Hi,

I'm using jax 0.1.35 on python 3.6.8. When evaluating gradients, I sometimes get `nan`s where I don't think `nan`s should be. This happens for some input values and not for others:

```
import jax
import jax.numpy as np
from jax import grad
print(jax.version.__version__)
x = np.linspace(190, 200, 4)
f = grad(lambda x: np.sum(np.tanh(x)))
print(x)
f(x)
```

```
0.1.35
[190.         193.33333333 196.66666667 200.        ]
DeviceArray([nan, nan, nan, nan], dtype=float32)
```

When tracing the `nan`s with `jax_debug_nans`, I get errors about divisions:

```
---------------------------------------------------------------------------
FloatingPointError                        Traceback (most recent call last)
<ipython-input-3-c4aba5ddce5a> in <module>
      6 f = grad(lambda x: np.sum(np.tanh(x)))
      7 print(x)
----> 8 f(x)

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/api.py in grad_f(*args, **kwargs)
    233   def grad_f(*args, **kwargs):
    234     if not has_aux:
--> 235       _, g = value_and_grad_f(*args, **kwargs)
    236       return g
    237     else:

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    287              "differentiation, pass holomorphic=True.")
    288       raise TypeError(msg.format(dtype))
--> 289     g = vjp_py(onp.ones((), dtype=dtype))
    290     g = g[0] if isinstance(argnums, int) else g
    291     if not has_aux:

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/api_util.py in apply_jaxtree_fun(fun, io_tree, *py_args)
     60       raise TypeError("Expected {}, got {}".format(expected, in_tree))
     61 
---> 62   ans = fun(*args)
     63   return build_tree(out_tree, ans)
     64 

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/api.py in out_vjp_packed(cotangent_in)
    820   ct_out_tree = PyTreeDef(node_types[tuple], None, in_trees)
    821   def out_vjp_packed(cotangent_in):
--> 822     return out_vjp(cotangent_in)
    823   vjp_py = partial(apply_jaxtree_fun, out_vjp_packed, (ct_in_trees, ct_out_tree))
    824   if not has_aux:

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/ad.py in vjp_(ct)
    110     dummy_primal_and_ct = pack((core.unit, ct))
    111     dummy_args = (None,) * len(jaxpr.invars)
--> 112     _, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primal_and_ct)
    113     return instantiate_zeros(pack(primals), arg_cts[1])
    114 

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, freevar_vals, args, cotangent_in)
    184       map(write_cotangent, bound_vars, ct_free_vars_out)
    185     else:
--> 186       cts_out = get_primitive_transpose(eqn.primitive)(ct_in, *invals, **eqn.params)
    187 
    188     if cts_out is zero:

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/lax/lax.py in _div_transpose_rule(cotangent, x, y)
   1565 def _div_transpose_rule(cotangent, x, y):
   1566   assert x is None and y is not None
-> 1567   res = ad_util.zero if cotangent is ad_util.zero else div(cotangent, y)
   1568   return res, None
   1569 div_p = standard_binop([_num, _num], 'div')

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/lax/lax.py in div(x, y)
    236 def div(x, y):
    237   r"""Elementwise division: :math:`x \\over y`."""
--> 238   return div_p.bind(x, y)
    239 
    240 def rem(x, y):

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/core.py in bind(self, *args, **kwargs)
    115     top_trace = find_top_trace(args)
    116     if top_trace is None:
--> 117       return self.impl(*args, **kwargs)
    118 
    119     tracers = map(top_trace.full_raise, args)

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **kwargs)
     50   abstract_args = map(abstractify, args)
     51   compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
---> 52   return compiled_fun(*args)
     53 
     54 @memoize

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/xla.py in execute_compiled_primitive(name, compiled, result_handler, *args)
     83   input_bufs = [device_put(x) for x in args]
     84   out_buf = compiled.Execute(input_bufs)
---> 85   check_nans(name, out_buf)
     86   return result_handler(out_buf)
     87 

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/xla.py in check_nans(name, buf)
     87 
     88 def check_nans(name, buf):
---> 89   FLAGS.jax_debug_nans and _check_nans(name, buf.shape(), buf)
     90 
     91 def _check_nans(name, xla_shape, buf):

/anaconda3/envs/py36_tf/lib/python3.6/site-packages/jax/interpreters/xla.py in _check_nans(name, xla_shape, buf)
     97       if onp.any(onp.isnan(pyval)):
     98         msg = "invalid value (nan) encountered in {}"
---> 99         raise FloatingPointError(msg.format(name))
    100 
    101 def device_put(x, device_num=0):

FloatingPointError: invalid value (nan) encountered in div
```

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

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 5feaa7924..442300300 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1437,7 +1437,7 @@ log1p_p = standard_unop(_float | _complex, 'log1p')
 ad.defjvp(log1p_p, lambda g, x: div(g, add(x, _one(x))))
 
 tanh_p = standard_unop(_float | _complex, 'tanh')
-ad.defjvp(tanh_p, lambda g, x: div(g, pow(cosh(x), _two(x))))
+ad.defjvp2(tanh_p, lambda g, ans, x: mul(g, sub(_one(x), mul(ans, ans))))
 
 sin_p = standard_unop(_float | _complex, 'sin')
 ad.defjvp(sin_p, lambda g, x: mul(g, cos(x)))
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index e04a16b3d..490721131 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1544,6 +1544,14 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
   def testIssue746(self):
     lnp.arange(12).reshape(3, 4)  # doesn't crash
 
+  def testIssue764(self):
+    x = lnp.linspace(190, 200, 4)
+    f = api.grad(lambda x: lnp.sum(lnp.tanh(x)))
+    # Expected values computed with autograd in float64 precision.
+    expected = onp.array([3.71669453e-165, 4.72999108e-168, 6.01954653e-171,
+                          7.66067839e-174], onp.float64)
+    self.assertAllClose(f(x), expected, check_dtypes=False)
+
 
 if __name__ == "__main__":
   absltest.main()