Repository: google/jax

Issue:
logsumexp: "invalid value (nan) encountered in convert_element_type"
I get this following error in logsumexp:
```
FloatingPointError: invalid value (nan) encountered in convert_element_type
/home/foo/miniconda3/envs/sica/lib/python3.7/site-packages/jax/_src/scipy/special.py(126)logsumexp()
```
Unfortunately, I have not been able to reproduce this in any reasonably sized version -- I only seem to encounter it within a much larger jitted and scanned code that I am not able to share at the moment due to anonymity requirements (earlier version of the code currently under review).  I tried to figure out what's going on by debugging the logsumexp function and found this strange:

```python
> out = lax.add(lax.log(jnp.sum(lax.exp(lax.sub(a, amax_with_dims))
> (Pdb) out
> Traced<ShapedArray(float64[2])>with<BatchTrace(level=1/0)>
>   with val = DeviceArray([[-0.8201989 , -0.8201989 ],
>                           [-0.82572116, -0.82572116],
>                           [-0.86567182, -0.86567182]], dtype=float64)
>        batch_dim = 0
> (Pdb) jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
>*** FloatingPointError: invalid value (nan) encountered in convert_element_type
> (Pdb) jnp.where(jnp.isnan(out), 0.0, 1.0).astype(out.dtype)
```
Yet
```python
>jnp.where(jnp.isnan(out), 0.0, 1.0).astype(out.dtype)
> Traced<ShapedArray(float64[2])>with<BatchTrace(level=1/0)>
>   with val = DeviceArray([[1., 1.],
>                           [1., 1.],
>                           [1., 1.]], dtype=float64)
>        batch_dim = 0
```
What makes this more strange, is that this only happens sometimes. e.g. if `out` is 2d rather than 3d then in never get this error.
Any ideas as to what might be going on? I will try to produce a minimal example somehow. 


Discussion:
Thanks for the report. One point of clarification: is this with `jax_debug_nans` enabled?
It is indeed.
(the issue happened originally without jax_debug_nans). However, I just updated JAX&JAXlib in hopes that might help but am now instead getting 
> Segmentation fault (core dumped)

So seems like there is maybe something else going on as well -- will try to fix that first and then see if this issue still persists.
So what seems to have happened:  1. there was a NaN coming from somewhere 2. I turned on jax_debug_nans in attempt to find out where it's coming from 3. This pointed to the logsumexp function 4. As per your comments, I turned off jax_debug_nans and debugged manually (by first disabling jit) 5. I found out the NaN was coming from matrix inversion -- completely different place than where jax_debug_nans pointed to...
Ah, thanks for the update. Sorry for the suboptimal behavior in `debug_nans`. It sounds like an issue we should investigate – is it possible to reproduce the issue compactly?
I created a much simplified version of my code that captures similar behaviour. If I run this as it is, it points to the error being in logsumexp but if I comment out with` jax.disable_jit():` then it appropriately shows the error being in Cholesky. Now I suppose the issue then is the interplay between debug_nans=True and disable_jit(), and a much simpler example would have sufficed -- should I not be using these two simultaneously? Reason why I did is that often in my code that has loads of scans and nested functions, in the past, I have only found debug_nans useful if jit is disabled...

```
import jax
import jax.numpy as jnp
from jax import jit
import pdb
from jax.lax import scan
from jax.scipy.special import logsumexp
from jax.config import config
config.update("jax_debug_nans", True)


# inv(L*L.T)*Y
def invcholp(L, Y):
    D = jax.scipy.linalg.solve_triangular(L, Y, lower=True)
    B = jax.scipy.linalg.solve_triangular(L.T, D, lower=False)
    return B


# inv(X)*Y
def invmp(X, Y):
    return invcholp(jnp.linalg.cholesky(X), Y)


def tree_prepend(prep, tree):
    preprended = jax.tree_multimap(
        lambda a, b: jnp.vstack((a[None], b)), prep, tree
    )
    return preprended


def scan_body(msg_in, msg_current):
    eta_A, eta_B = msg_current
    fwd_t = logsumexp(eta_A.T + msg_in[None, :], axis=1)
    msg_out = fwd_t+eta_B
    return msg_out, fwd_t


def tree_stuff(eta_pi, eta_A):
    msg_in = jnp.array([0.04, 0.98])
    eta_B = jnp.log(eta_A + 1000)
    tree = tree_prepend(eta_pi, scan(scan_body, msg_in, (eta_A, eta_B))[1])
    return tree


@jit
def run(V, eta_pi, eta_A):
    eta_A = jnp.repeat(eta_A, 1000, 0)
    tree_stuff(eta_pi, eta_A)
    Q = invmp(V, jnp.eye(V.shape[0]))
    return Q


if __name__ == "__main__":
    V = jnp.array([[-10, 20], [-10, 20]])
    eta_pi = jnp.array([0.1])
    eta_A = jnp.array([[0.5, 0.08],
                       [0.08, 0.3]])
    with jax.disable_jit():
        run(V, eta_pi, eta_A)
```
Also, if instead of  `jax.disable_jit()` I comment out the jit decorator of the `run` function, as a way to disable jit, then behaviour is as expected. 
Are you using complex numbers by any chance? I have somewhat recently made a PR (#6855) to support complex numbers in `logsumexp`. The code path used for `complex64/128` is different by the one used for the other dtypes.
Thanks for the repro: I think the issue is that `tree_stuff(eta_pi, eta_A)` is dead code (it returns a value that is unused within the function). One optimization done by `jit` is removal of dead code.

So, without `jit` this dead code executes and gives an error. With `jit`, the dead code is removed and a different error comes from the next line, `Q = invmp(V, jnp.eye(V.shape[0]))`

This appears to all be working as expected.
> Are you using complex numbers by any chance? I have somewhat recently made a PR (#6855) to support complex numbers in `logsumexp`. The code path used for `complex64/128` is different by the one used for the other dtypes.

No complex numbers.

> Thanks for the repro: I think the issue is that tree_stuff(eta_pi, eta_A) is dead code (it returns a value that is unused within the function). One optimization done by jit is removal of dead code.
> 
> So, without jit this dead code executes and gives an error. With jit, the dead code is removed and a different error comes from the next line, Q = invmp(V, jnp.eye(V.shape[0]))
> 
> This appears to all be working as expected.

Hmm interesting. This does sounds like it makes sense here, but in my original code the output of `tree_stuff` is definitely used within `run`. 
Thanks - trying to simplify things a bit I do find something strange:
```python
import jax
import jax.numpy as jnp
from jax import jit
from jax.lax import scan
from jax.scipy.special import logsumexp
from jax.config import config
config.update("jax_debug_nans", True)

def scan_body(msg_in, msg_current):
    eta_A, eta_B = msg_current
    fwd_t = logsumexp(eta_A + msg_in, axis=0)
    return fwd_t + eta_B, fwd_t

@jit
def run(V, eta_pi, eta_A):
    eta_A = jnp.repeat(eta_A, 1000, 0)
    msg_in = jnp.array([0.04, 0.98])
    eta_B = jnp.log(eta_A + 1000)
    return scan(scan_body, msg_in, (eta_A, eta_B))[1]

V = jnp.array([[-10, 20], [-10, 20]])
eta_pi = jnp.array([0.1])
eta_A = jnp.array([[0.5, 0.08],
                    [0.08, 0.3]])

# runs without issue
print(run(V, eta_pi, eta_A))

# runs without issue
print(jit(run)(V, eta_pi, eta_A))

# raises an error due to NaN in convert_element_type
with jax.disable_jit():
    print(jit(run)(V, eta_pi, eta_A))
```
When the function is run without jit or with jit, it succeeds. But when it is jitted and jit is disabled, it fails. I'm not sure what is going on with that.
Yes I only saw issues whenever I had `with jax.disable_jit():`. Disabling all jit decorators manually didnt produce any issues. 
> Thanks - trying to simplify things a bit I do find something strange:
> 
> ```python
> import jax
> import jax.numpy as jnp
> from jax import jit
> from jax.lax import scan
> from jax.scipy.special import logsumexp
> from jax.config import config
> config.update("jax_debug_nans", True)
> 
> def scan_body(msg_in, msg_current):
>     eta_A, eta_B = msg_current
>     fwd_t = logsumexp(eta_A + msg_in, axis=0)
>     return fwd_t + eta_B, fwd_t
> 
> @jit
> def run(V, eta_pi, eta_A):
>     eta_A = jnp.repeat(eta_A, 1000, 0)
>     msg_in = jnp.array([0.04, 0.98])
>     eta_B = jnp.log(eta_A + 1000)
>     return scan(scan_body, msg_in, (eta_A, eta_B))[1]
> 
> V = jnp.array([[-10, 20], [-10, 20]])
> eta_pi = jnp.array([0.1])
> eta_A = jnp.array([[0.5, 0.08],
>                     [0.08, 0.3]])
> 
> # runs without issue
> print(run(V, eta_pi, eta_A))
> 
> # runs without issue
> print(jit(run)(V, eta_pi, eta_A))
> 
> # raises an error due to NaN in convert_element_type
> with jax.disable_jit():
>     print(jit(run)(V, eta_pi, eta_A))
> ```
> 
> When the function is run without jit or with jit, it succeeds. But when it is jitted and jit is disabled, it fails. I'm not sure what is going on with that.

(though when I run your simplified code I don't see any errors, perhaps because the matrix V is not inverted anywhere)
> (though when I run your simplified code I don't see any errors, perhaps because the matrix V is not inverted anywhere)

That's my point: the error does not seem related to a matrix inversion, but rather to the interaction between `scan` and `disable_jit` (note that even in your original code, the `logsumexp` error is raised before any matrix inversion takes place).

The simplified code as written errors when I run it on CPU with JAX v0.2.19. What version of JAX are you using, on which backend?
Apologies, I had not activated my proper jax conda env so it was running it with much older 0.1.55 . Switching to CPU JAX 0.2.19 errors our precisely as you describe (unfortunately I dont have GPU access at the moment so can't check if that would change things).
Here's an even more stripped-down version of the error, showing that `debug_nans` triggers with only if `scan` is run under `disable_jit`:

```python
import jax
from jax.lax import scan
import jax.numpy as jnp
from jax.scipy.special import logsumexp
from jax.config import config
config.update("jax_debug_nans", True)

def scan_body(carry, x):
  return logsumexp(x), x

carry = 0.0
x = jnp.ones(2)

print(scan(scan_body, carry, x))
# (DeviceArray(1., dtype=float32), DeviceArray([1., 1.], dtype=float32))

with jax.disable_jit():
  print(scan(scan_body, carry, x))
  # FloatingPointError: invalid value (nan) encountered in convert_element_type
```
This seems to be a consequence of the alternate implementation of `scan` under `disable_jit`: https://github.com/google/jax/blob/476642578be3092db77adf0bdeca5cd4867f12e1/jax/_src/lax/control_flow.py#L1252-L1263

Though I'm not entirely certain why this would be different in terms of generating nans.
Why would it be that changing the "official" logsumexp code from
`sign = jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
`
to something else  like
`sign = jnp.where(jnp.isnan(out), 0., 1.0).astype(out.dtype)`
doesnt cause error in your new simplified example. Is there some problem with using ´np.nan¨ from original numy?

Fix is in #7648.

The issue is that if you use `np.nan` in `_where()` with `disable_jit` and `debug_nans`, it triggers a NaN error because the value is treated as a constant rather than being traced. We can work around it by explicitly casting the value to a JAX type with `jnp.array`
Awesome that makes sense -- thanks very much! 

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
@@ -120,8 +120,8 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     out = lax.add(lax.log(jnp.sum(lax.exp(lax.sub(a, amax_with_dims)),
                                   axis=dims, keepdims=keepdims)),
                   amax)
-    sign = jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
-    sign = jnp.where(jnp.isneginf(out), 0.0, sign)
+    sign = jnp.where(jnp.isnan(out), out, 1.0)
+    sign = jnp.where(jnp.isneginf(out), 0.0, sign).astype(out.dtype)
   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
@@ -214,6 +214,13 @@ def testLogSumExpOnes(self):
       self._CheckAgainstNumpy(osp_special.logsumexp, lsp_special.logsumexp, args_maker)
       self._CompileAndCheck(lsp_special.logsumexp, args_maker)
 
+  def testLogSumExpNans(self):
+    # Regression test for https://github.com/google/jax/issues/7634
+    with jax.debug_nans(True):
+      with jax.disable_jit():
+        result = lsp_special.logsumexp(1.0)
+    self.assertEqual(result, 1.0)
+
   @parameterized.named_parameters(itertools.chain.from_iterable(
     jtu.cases_from_list(
         {"testcase_name": jtu.format_test_name_suffix(