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?



Discussion:
Thanks for the report - I'm not sure why this changed between versions, but I think I know the fix. I'll try to get it in today

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
@@ -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:


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
@@ -25,6 +25,7 @@
 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 @@ def testLogSumExpZeros(self):
     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(