Repository: google/jax

Issue:
jax.scipy.special.logsumexp raises OverflowError
`jax.scipy.special.logsumexp` seems to be broken in jax 0.2.9.

The following code:
```
import jax, scipy
print(jax.scipy.special.logsumexp(jax.numpy.array([-1, 1])))
print(scipy.special.logsumexp(jax.numpy.array([-1, 1])))
```

Raises the following error:
```Traceback (most recent call last):
  File "txt.py", line 2, in <module>
    print(jax.scipy.special.logsumexp(jax.numpy.array([-1, 1])))
  File "/opt/conda/lib/python3.7/site-packages/jax/_src/scipy/special.py", line 107, in logsumexp
    amax = lax.reduce(a, _constant_like(a, -np.inf), lax.max, dims)
  File "/opt/conda/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py", line 335, in _constant_like
    return np.array(const, dtype=_dtype(x))
OverflowError: cannot convert float infinity to integer```

While the line from scipy runs fine. This was reproduced for the handful of other inputs I tried


Discussion:
`jax.scipy.special.logsumexp` does not work on integers. If you use floats like

```python
jax.scipy.special.logsumexp(jax.numpy.array([-1., 1.]))
# -> DeviceArray(1.1269281, dtype=float32)
```

it works. You may still consider this a bug, just FYI.
Thanks for the report. Indeed I think the issue is that `logsumexp` currently fails for integer input (although the error is more explicit on master due to changes in #5872). The fix will be to internally promote inputs to floating point; the workaround right now is to convert inputs manually:
```python
import jax, scipy
print(jax.scipy.special.logsumexp(jax.numpy.array([-1, 1]).astype(jnp.float32)))
```

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -103,6 +103,8 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   if b is not None:
     a, b = _promote_args_inexact("logsumexp", a, b)
     a = jnp.where(b != 0, a, -jnp.inf)
+  else:
+    a, = _promote_args_inexact("logsumexp", a)
   pos_dims, dims = _reduction_dims(a, axis)
   amax = jnp.max(a, axis=dims, keepdims=keepdims)
   amax = lax.stop_gradient(lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0)))


Test patch:
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
@@ -106,7 +106,7 @@ def _GetArgsMaker(self, rng, shapes, dtypes):
        "shapes": shapes, "dtype": dtype,
        "axis": axis, "keepdims": keepdims,
        "return_sign": return_sign, "use_b": use_b}
-      for shape_group in compatible_shapes for dtype in float_dtypes
+      for shape_group in compatible_shapes for dtype in float_dtypes + int_dtypes
       for use_b in [False, True]
       for shapes in itertools.product(*(
         (shape_group, shape_group) if use_b else (shape_group,)))
@@ -143,8 +143,9 @@ def lax_fun(array_to_reduce):
                                      return_sign=return_sign)
 
       args_maker = lambda: [rng(shapes[0], dtype)]
+    tol = {np.float32: 1E-6, np.float64: 1E-14}
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
-    self._CompileAndCheck(lax_fun, args_maker)
+    self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
 
   def testLogSumExpZeros(self):
     # Regression test for https://github.com/google/jax/issues/5370