Repository: google/jax
Issue: Debug infs and logsumexp

I've found a difference in behavior between the latest versions and a (recent) previous version of the library I was using. I found this trying to debug a larger issue on an hpc application being build using jax but when trying to debug that, this was thrown which might point to the root cause of the larger issue.

I wanted to clarify if the following is expected:

Specifically it appears that when trying to use the debug_infs config option, with a special function that does some checking for np.inf itself (in this case logsumexp), the latest version of the library from pip will throw an error that was passed previously

MWE:
```python
from jax.scipy.special import logsumexp

from jax.config import config
config.update("jax_debug_infs", True)

import numpy as np

x=np.ones(10)
y=logsumexp(x)

print(y)
```

Using jaxlib 0.1.64 and jax 0.2.10 this prints the value fine (note none of the values in the array that is being summed are inf)


updating to jaxlib 0.1.69 and jax 0.2.18 the following error shows:

```
Traceback (most recent call last):
  File "test_inf.py", line 9, in <module>
    y=logsumexp(x)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/scipy/special.py", line 124, in logsumexp
    sign = jnp.where(out == -np.inf, 0.0, sign)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5871, in deferring_binary_op
    return binary_op(self, other)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 421, in <lambda>
    fn = lambda x1, x2: lax_fn(*_promote_args(numpy_fn.__name__, x1, x2))
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 326, in _promote_args
    return _promote_shapes(fun_name, *_promote_dtypes(*args))
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 272, in _promote_dtypes
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 272, in <listcomp>
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 454, in _convert_element_type
    return convert_element_type_p.bind(operand, new_dtype=new_dtype,
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/core.py", line 264, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/core.py", line 604, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/interpreters/xla.py", line 262, in apply_primitive
    return compiled_fun(*args)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/interpreters/xla.py", line 379, in _execute_compiled_primitive
    check_special(prim.name, out_bufs)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/interpreters/xla.py", line 398, in check_special
    _check_special(name, buf.xla_shape(), buf)
  File "/rds/user/dy297/hpc-work/install/python/lib/python3.8/site-packages/jax/interpreters/xla.py", line 406, in _check_special
    raise FloatingPointError(f"invalid value (inf) encountered in {name}")
FloatingPointError: invalid value (inf) encountered in convert_element_type

```

And it appears to be pulling up the -np.inf used to check inside `jnp.where(out == -np.inf, 0.0, sign)`

Is it expected that these sorts of condition checks inside special functions should fail the debug configs?


Changed files:
tests/lax_scipy_test.py, jax/_src/scipy/special.py

diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index 7c39db27f..546a2ae05 100644
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -121,7 +121,7 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
                                   axis=dims, keepdims=keepdims)),
                   amax)
     sign = jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
-    sign = jnp.where(out == -np.inf, 0.0, sign)
+    sign = jnp.where(jnp.isneginf(out), 0.0, sign)
   else:
     expsub = lax.exp(lax.sub(a, amax_with_dims))
     if b is not None:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index 25c5215a3..67ef1b2db 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -25,6 +25,7 @@ from absl.testing import parameterized
 import numpy as np
 import scipy.special as osp_special
 
+import jax
 from jax._src import api
 from jax import numpy as jnp
 from jax import lax
@@ -205,6 +206,13 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
     self._CompileAndCheck(lax_fun, args_maker)
 
+  def testLogSumExpOnes(self):
+    # Regression test for https://github.com/google/jax/issues/7390
+    args_maker = lambda: [np.ones(4, dtype='float32')]
+    with jax.debug_infs(True):
+      self._CheckAgainstNumpy(osp_special.logsumexp, lsp_special.logsumexp, args_maker)
+      self._CompileAndCheck(lsp_special.logsumexp, args_maker)
+
   @parameterized.named_parameters(itertools.chain.from_iterable(
     jtu.cases_from_list(
         {"testcase_name": jtu.format_test_name_suffix(