Repository: google/jax

Issue:
logaddexp and logsumexp do not handle -inf correctly, don't match numpy/scipy
logaddexp(-inf, -inf) should be -inf, is nan in jax implementations.

numpy.logaddexp(np.NINF, np.NINF) => -inf
jax.numpy.logaddexp(np.NINF, np.NINF) => nan

Same issue for special.logsumexp


Discussion:
I'm seeing -inf for logaddexp but NaN for logsumexp. Can you share the exact code that gives you NaN with jax.numpy.logaddexp?
The jax.numpy.logaddexp behavior is really weird, I only see it when running on an array with more than one item:

import jax.numpy as np
np.logaddexp(np.NINF, np.NINF) -> -inf
np.logaddexp(np.array([np.NINF]), np.array([np.NINF])) -> [-inf]
np.logaddexp(np.array([np.NINF, -1.0]), np.array([np.NINF, -1.0])) -> [nan, -0.30685282]
I think the last behavior is because of XLA mishandling `exp` for NaN inputs on CPU:

```
In [7]: lax.exp(np.nan)
Out[7]: DeviceArray(0., dtype=float32)

In [8]: lax.exp(np.array([np.nan, 0.]))
Out[8]: DeviceArray([inf,  1.], dtype=float32)
```

I filed a bug with the XLA folks.

Resolving patch:
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
@@ -516,7 +516,10 @@ def logaddexp(x1, x2):
   x1, x2 = _promote_shapes("logaddexp",
                            *_promote_to_result_dtype(onp.logaddexp, x1, x2))
   amax = lax.max(x1, x2)
-  return lax.add(amax, lax.log1p(lax.exp(-lax.abs(lax.sub(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)))))
 
 
 @_wraps(onp.logaddexp2)
@@ -524,8 +527,12 @@ def logaddexp2(x1, x2):
   x1, x2 = _promote_shapes("logaddexp2",
                            *_promote_to_result_dtype(onp.logaddexp2, x1, x2))
   amax = lax.max(x1, x2)
-  return lax.add(amax, lax.div(lax.log1p(exp2(-lax.abs(lax.sub(x1, x2)))),
-                               _constant_like(x1, onp.log(2))))
+  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.div(lax.log1p(exp2(-lax.abs(delta))),
+                                          _constant_like(x1, onp.log(2)))))
+
 
 @_wraps(onp.log2)
 def log2(x):
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -81,6 +81,7 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   shape = lax.subvals(onp.shape(a), zip(dims, (1,) * len(dims)))
   dimadd = lambda x: lax.reshape(x, shape)
   amax = lax.reduce(a, _constant_like(a, -onp.inf), lax.max, dims)
+  amax = lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0))
   amax_singletons = dimadd(amax)
   out = lax.add(lax.log(lax.reduce(lax.exp(lax.sub(a, amax_singletons)),
                                    _constant_like(a, 0), lax.add, dims)), amax)


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
@@ -191,9 +191,11 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None,
               test_name="log1p_large", tolerance={onp.float64: 1e-12}),
     op_record("log1p", 1, number_dtypes, all_shapes, jtu.rand_small_positive(), [],
               tolerance={onp.float64: 1e-12}),
-    op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"],
+    op_record("logaddexp", 2, float_dtypes, all_shapes,
+              jtu.rand_some_inf_and_nan(), ["rev"],
               tolerance={onp.float64: 1e-12}),
-    op_record("logaddexp2", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"],
+    op_record("logaddexp2", 2, float_dtypes, all_shapes,
+              jtu.rand_some_inf_and_nan(), ["rev"],
               tolerance={onp.float16: 1e-2}),
     op_record("polyval", 2, number_dtypes, nonempty_nonscalar_array_shapes,
               jtu.rand_default(), [], check_dtypes=False,
@@ -1851,7 +1853,7 @@ def testMathSpecialFloatValues(self, op, dtype):
     onp_op = getattr(onp, op)
     lnp_op = getattr(lnp, op)
     dtype = onp.dtype(xla_bridge.canonicalize_dtype(dtype)).type
-    for x in (onp.nan, -onp.inf, -100., -2. -1., 0., 1., 2., 100., onp.inf,
+    for x in (onp.nan, -onp.inf, -100., -2., -1., 0., 1., 2., 100., onp.inf,
               lnp.finfo(dtype).max, onp.sqrt(lnp.finfo(dtype).max),
               onp.sqrt(lnp.finfo(dtype).max) * 2.):
       if onp.isnan(x) and op in ("sinh", "cosh", "expm1", "exp"):
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -84,7 +84,10 @@ def _GetArgsMaker(self, rng, shapes, dtypes):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_inshape={}_axis={}_keepdims={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis, keepdims),
-       "rng": jtu.rand_default(), "shape": shape, "dtype": dtype,
+       # TODO(b/133842870): re-enable when exp(nan) returns NaN on CPU.
+       "rng": jtu.rand_some_inf_and_nan() if jtu.device_under_test() != "cpu"
+              else jtu.rand_default(),
+       "shape": shape, "dtype": dtype,
        "axis": axis, "keepdims": keepdims}
       for shape in all_shapes for dtype in float_dtypes
       for axis in range(-len(shape), len(shape))