Repository: google/jax

Issue:
grad(softplus) has wrong value at zero
`jax.grad(jax.nn.softplus)(0.0)` evaluates to 0.0, which is definitely wrong -- the right answer is 0.5.

This is easy to visualize:
```python
import matplotlib.pyplot as plt
x = jax.numpy.linspace(-2, 2, num=101)
plt.plot(x, jax.vmap(jax.grad(jax.nn.softplus))(x))
```
![image](https://user-images.githubusercontent.com/1217238/73335588-b4ee6680-4224-11ea-8061-3982ea777051.png)



Discussion:
`logaddexp` is implemented as:
```
  amax = lax.max(x1, x2)
  delta = lax.sub(x1, x2)
  return lax.select(isnan(delta),
                    lax.add(x1, x2),  # NaNs or infinities of the same sign.
                    lax.add(amax, lax.log1p(lax.exp(-lax.abs(delta)))))
```
The derivative of the implementation, especially the `abs` and `max`, seems problematic.

We'd probably do best to either make `logaddexp` primitive or to add a custom JVP rule for it.
We have run into this on a real project so I think it would be good if we could fix this as soon as possible.

How about this version?
```python
def logaddexp2(x1, x2):
  amax = lax.max(x1, x2)
  amin = lax.min(x1, x2)
  delta = amin - amax
  return lax.select(jnp.isnan(delta),
                    lax.add(x1, x2),
                    lax.add(amax, lax.log1p(lax.exp(delta)))
```
It seems to be slightly slower than the original:
``` python
xs = np.linspace(-5, 5, 1000001)
f1 = jax.jit(jax.vmap(jax.grad(logaddexp)))
f2 = jax.jit(jax.vmap(jax.grad(logaddexp2)))
ys = np.zeros_like(xs)
f1(xs, ys)
f2(xs, ys)
%timeit f1(xs, ys)
%timeit f2(xs, ys)
```
```
100 loops, best of 3: 5.79 ms per loop
100 loops, best of 3: 6.35 ms per loop
```
But the slowdown is not too bad, should I make a PR with this version as an interim solution?

`0` seems to be handled correctly with this:
![image](https://user-images.githubusercontent.com/890531/74037740-d08d0600-49b6-11ea-98a3-e018cee438c0.png)
I'm attempting a solution that changes lax.max grad, #2195 
@ibab That seems reasonable, though I'm not sure we can rely on XLA preserving NaN semantics under a `jit` on all backends.

@joaogui1 Thanks for digging in! I'm not sure we should couple fixing this issue to changing our derivative for `lax.max`, though. We can discuss what the best policy for `lax.max` is, but this issue seems to be crying out for more direct solution, like a custom JVP rule.
@ibab re: NaN semantics, I have learned that XLA _will_ respect NaN semantics (and while LLVM on CPU might not, JAX currently sets flags so that it does). So your solution seems viable.

We're looking into this more...
@mattjj okss, but I think we should take a look at `lax.max` later, as right now it's gradient is not compatible with `lax.abs` gradient

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -143,8 +143,11 @@ def backward_pass(jaxpr: core.Jaxpr, consts, args, cotangents_in):
 
   def write_cotangent(v, ct):
     # assert v not in primal_env
-    if ct is not None and type(v) is not Literal:
+    if ct is not None and type(v) is not Literal and ct is not zero:
       ct_env[v] = add_tangents(ct_env[v], ct) if v in ct_env else ct
+      if not core.skip_checks:
+        ct_aval = core.get_aval(ct_env[v])
+        assert v.aval == core.lattice_join(v.aval, ct_aval)
 
   def read_cotangent(v):
     return ct_env.get(v, zero)
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -20,7 +20,7 @@
                   _const, _eq_meet, _broadcasting_select,
                   _check_user_dtype_supported, _one, _const,
                   _upcast_fp16_for_computation, _broadcasting_shape_rule,
-                  _eye, _tri, _delta)
+                  _eye, _tri, _delta, _ones, _zeros)
 from .lax_control_flow import *
 from .lax_fft import *
 from .lax_parallel import *
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -39,7 +39,7 @@
 import numpy as onp
 import opt_einsum
 
-from jax import jit, device_put
+from jax import jit, device_put, custom_jvp
 from .. import core
 from .. import dtypes
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray
@@ -583,6 +583,7 @@ def power(x1, x2):
   return acc
 
 
+@custom_jvp
 @_wraps(onp.logaddexp)
 def logaddexp(x1, x2):
   x1, x2 = _promote_shapes("logaddexp", *_promote_dtypes_inexact(x1, x2))
@@ -592,7 +593,21 @@ def logaddexp(x1, x2):
                     lax.add(x1, x2),  # NaNs or infinities of the same sign.
                     lax.add(amax, lax.log1p(lax.exp(-lax.abs(delta)))))
 
+@logaddexp.defjvp
+def _logaddexp_jvp(primals, tangents):
+  x1, x2 = primals
+  t1, t2 = tangents
+  x1, x2, t1, t2 = broadcast_arrays(x1, x2, t1, t2)
+  primal_out = logaddexp(x1, x2)
+  tangent_out = (t1 * exp(_replace_inf(x1) - _replace_inf(primal_out)) +
+                 t2 * exp(_replace_inf(x2) - _replace_inf(primal_out)))
+  return primal_out, tangent_out
 
+def _replace_inf(x):
+  return lax.select(isposinf(x), zeros_like(x), x)
+
+
+@custom_jvp
 @_wraps(onp.logaddexp2)
 def logaddexp2(x1, x2):
   x1, x2 = _promote_shapes("logaddexp2", *_promote_dtypes_inexact(x1, x2))
@@ -602,6 +617,15 @@ def logaddexp2(x1, x2):
                     lax.add(x1, x2),  # NaNs or infinities of the same sign.
                     lax.add(amax, lax.div(lax.log1p(exp2(-lax.abs(delta))),
                                           _constant_like(x1, onp.log(2)))))
+@logaddexp2.defjvp
+def _logaddexp2_jvp(primals, tangents):
+  x1, x2 = primals
+  t1, t2 = tangents
+  x1, x2, t1, t2 = broadcast_arrays(x1, x2, t1, t2)
+  primal_out = logaddexp2(x1, x2)
+  tangent_out = (t1 * 2 ** (_replace_inf(x1) - _replace_inf(primal_out)) +
+                 t2 * 2 ** (_replace_inf(x2) - _replace_inf(primal_out)))
+  return primal_out, tangent_out
 
 
 @_wraps(onp.log2)


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -696,8 +696,10 @@ def cases_from_gens(*gens):
 class JaxTestCase(parameterized.TestCase):
   """Base class for JAX tests including numerical checks and boilerplate."""
 
-  def tearDown(self) -> None:
-    assert core.reset_trace_state()
+  # TODO(mattjj): this obscures the error messages from failures, figure out how
+  # to re-enable it
+  # def tearDown(self) -> None:
+  #   assert core.reset_trace_state()
 
   def assertArraysAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
     """Assert that x and y are close (up to numerical tolerances)."""
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
@@ -2966,6 +2966,12 @@ def grad_test_spec(op, nargs, order, rng_factory, dtypes, name=None, tol=None):
     grad_test_spec(jnp.arctanh, nargs=1, order=2,
                    rng_factory=partial(jtu.rand_uniform, -0.9, 0.9),
                    dtypes=[onp.float64, onp.complex64], tol=1e-4),
+    grad_test_spec(jnp.logaddexp, nargs=2, order=1,
+                   rng_factory=partial(jtu.rand_uniform, -0.9, 0.9),
+                   dtypes=[onp.float64], tol=1e-4),
+    grad_test_spec(jnp.logaddexp2, nargs=2, order=2,
+                   rng_factory=partial(jtu.rand_uniform, -0.9, 0.9),
+                   dtypes=[onp.float64], tol=1e-4),
 ]
 
 GradSpecialValuesTestSpec = collections.namedtuple(
@@ -2975,7 +2981,7 @@ def grad_test_spec(op, nargs, order, rng_factory, dtypes, name=None, tol=None):
     GradSpecialValuesTestSpec(jnp.arcsinh, [0., 1000.], 2),
     GradSpecialValuesTestSpec(jnp.arccosh, [1000.], 2),
     GradSpecialValuesTestSpec(jnp.arctanh, [0.], 2),
-    GradSpecialValuesTestSpec(jnp.sinc, [0.], 1)
+    GradSpecialValuesTestSpec(jnp.sinc, [0.], 1),
 ]
 
 def num_float_bits(dtype):
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -39,6 +39,27 @@ def testSoftplusGrad(self):
     check_grads(nn.softplus, (1e-8,), order=4,
                 rtol=1e-2 if jtu.device_under_test() == "tpu" else None)
 
+  def testSoftplusGradZero(self):
+    check_grads(nn.softplus, (0.,), order=1,
+                rtol=1e-2 if jtu.device_under_test() == "tpu" else None)
+
+  def testSoftplusGradInf(self):
+    self.assertAllClose(
+        1., jax.grad(nn.softplus)(float('inf')), check_dtypes=True)
+
+  def testSoftplusGradNegInf(self):
+    check_grads(nn.softplus, (-float('inf'),), order=1,
+                rtol=1e-2 if jtu.device_under_test() == "tpu" else None)
+
+  def testSoftplusGradNan(self):
+    check_grads(nn.softplus, (float('nan'),), order=1,
+                rtol=1e-2 if jtu.device_under_test() == "tpu" else None)
+
+  @parameterized.parameters([
+      int, np.int32, float, np.float64, np.float32, np.float64,])
+  def testSoftplusZero(self, dtype):
+    self.assertEqual(np.log(dtype(2)), nn.softplus(dtype(0)))
+
   def testReluGrad(self):
     rtol = 1e-2 if jtu.device_under_test() == "tpu" else None
     check_grads(nn.relu, (1.,), order=3, rtol=rtol)