Repository: google/jax

Issue:
`log_softmax` and `softmax` are not numerically stable with `where` argument inside `vmap`
### Description

In the following code, large masked-out values seem to make the gradient infeasible. Uncommenting `x = jnp.where(mask, x, x.min())` fixes the issue. Apologies if I'm not using log_softmax right.

```python
import os
import jax
import jax.numpy as jnp
import jax.random as jr
from jax import vmap, jit, nn, value_and_grad

key = jr.PRNGKey(12322351)
@jit
def f(x, mask, i):
    # x = jnp.where(mask, x, x.min())
    return nn.log_softmax(x, where=mask, initial=x.max())[i]

key, sub = jr.split(key)
x = jr.normal(sub, (1000, 12)) * 100
key, sub = jr.split(key)
mask = jr.bernoulli(sub, 0.5, (1000, 12))
idx = jnp.zeros(1000, dtype=jnp.int32)
for i in range(1000):
    while not jnp.any(mask[i]):
        key, sub = jr.split(key)
        mask = mask.at[i].set(jr.bernoulli(sub, 0.5, (12,)))
    j, = jnp.where(mask[i])
    key, sub = jr.split(key)
    idx = idx.at[i].set(jr.choice(sub, j))

v, g = value_and_grad(lambda x: vmap(f)(x, mask, idx).sum())(x)
v.item(), g.min().item(), g.max().item()
```
returns
```
(inf, nan, nan)
```

### What jax/jaxlib version are you using?

jax==0.4.23
jaxlib==0.4.23+cuda12.cudnn89

### Which accelerator(s) are you using?

GPU

### Additional system info?

1.26.2 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:40:35) [GCC 12.3.0] uname_result(system='Linux', node='xxxx', release='5.14.0-284.30.1.el9_2.x86_64', version='#1 SMP PREEMPT_DYNAMIC Tue Sep 12 09:28:32 EDT 2023', machine='x86_64')

### NVIDIA GPU info

```
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.06              Driver Version: 545.23.06    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A100-SXM4-80GB          On  | 00000000:07:00.0 Off |                    0 |
| N/A   31C    P0              70W / 400W |    971MiB / 81920MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+

```


Discussion:
Thanks for the report – this looks like an instance of the issue discussed here: https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where

It seems like you arrived at the correct fix, which is applying `where` twice to avoid passing problematic values to the differentiated function. My only question here is whether `softmax` and other functions providing the `where` argument should do this double-where by default... I'm not sure. What do you think?
I personally did not think my nans were coming from the softmax, so I'd lean on doing the double where in softmax. Perhaps it could be disabled with a flag for the performance-conscious people?
Here's a much more compact repro:
```python
import jax
import jax.numpy as jnp

def f(x, mask):
    # x = jnp.where(mask, x, x.min())
    return jax.nn.log_softmax(x, where=mask, initial=x.max())[0]

x = jnp.array([36., 200.])
mask = jnp.array([True,  False])
print(*jax.value_and_grad(f)(x, mask))
# inf [nan  0.]
```
Oh, I see the issue now – it's a problem with the `initial` value. You're setting `initial=x.max()`, when you should set `initial=x.min()`.

The gradient will still contain `NaN` value because the output of `logsoftmax` contains infinities where `mask` is False. This is the expected output, because `softmax(x) = 0`, so `logsoftmax(x) = -inf`, and gradients of infinity are `NaN` because no other value makes sense.

If you don't want `NaN` gradients, you need to take the derivative of a function that does not have infinite outputs.

Does that make sense?
I don't think that's expected, though. I'm only looking at indexes where softmax is **not** zero. The same issue happens with softmax instead of logsoftmax where the gradient should be zero because the output is independent of the input. This only happens with the vmap. This also happens when initial x.min() or x.max()
I'm not sure I follow – could you put together a minimal example that demonstrates these issues? i.e. one that uses a constant array with a few values like my repro, rather than an expensive-to-construct array with thousands of entries as in your original repro.
Will do!
You were right that this does not seem to have anything to do with Vmap, and I might be misunderstanding what you are saying, but please look at the following piece of code.

```python
def f(x, mask):
    # x = jnp.where(mask, x, x.min())
    return jax.nn.log_softmax(x, where=mask, initial=x[jnp.argmax(mask)])[0]

mask = jnp.array([True,  False])
print(*jax.value_and_grad(f)(jnp.array([36., 36.]), mask))
print(*jax.value_and_grad(f)(jnp.array([36., 10000.]), mask))
# 0.0 [0. 0.]
# 0.0 [ 0. nan]
```
I think that the behavior should be consistent, either always nan or always zero. 

> The gradient will still contain NaN value because the output of logsoftmax contains infinities where mask is False.

I don't think that this statement is true. The gradient doesn't have nans yet there is a mask. The output does not depend on the -infs, so the gradient should still be defined. If a value is masked out, then the gradient should just be zero like it is here when both values are close. 

Thanks - that's very clear. I was confused by the claim that this was related to `vmap`, but this makes more sense.

Side-note: as noted above, you should set `initial` to the minimum value in the array (or negative infinity in case the mask is all `False`). That won't fix the `nan` issue, but it is what the function expects.

I think I was correct when I speculated above that the `nan` here is coming from the issue mentioned in the linked FAQ entry. The way to address it is to use the double-where construct, which can be accomplished by uncommenting your commented line. Perhaps `softmax` should do that automatically, but that's less clear to me.
The `nan` in the output is coming from the `sum` in the implementation; here's a minimization of the implementation issue:
```python
def f(x, mask):
  return jnp.sum(jnp.exp(x), where=mask)

mask = jnp.array([True, False])
print(*jax.value_and_grad(f)(jnp.array([0., 10000.]), mask))
# 1.0 [ 1. nan]
```
So perhaps `sum` should include the double-where construct to avoid potentially generating NaNs in gradients? If we fixed it in sum, it would be fixed in softmax as well.
That sounds reasonable to me
This should be fixed in #19518. I've confirmed that it fixes the initial report so long as you use `initial=x.min()` instead of `initial=x.max()`.

Resolving patch:
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -432,7 +432,8 @@ def log_softmax(x: ArrayLike,
   numpy_util.check_arraylike("log_softmax", x)
   x_arr = jnp.asarray(x)
   x_max = jnp.max(x_arr, axis, where=where, initial=initial, keepdims=True)
-  shifted = x_arr - lax.stop_gradient(x_max)
+  x_safe = x_arr if where is None else jnp.where(where, x_arr, initial)
+  shifted = x_safe - lax.stop_gradient(x_max)
   shifted_logsumexp = jnp.log(
       jnp.sum(jnp.exp(shifted), axis, where=where, keepdims=True))
   result = shifted - shifted_logsumexp
@@ -486,7 +487,8 @@ def _softmax(
     where: ArrayLike | None = None,
     initial: ArrayLike | None = None) -> Array:
   x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True)
-  unnormalized = jnp.exp(x - x_max)
+  x_safe = x if where is None else jnp.where(where, x, initial)
+  unnormalized = jnp.exp(x_safe - x_max)
   result = unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
   if where is not None:
     result = jnp.where(where, result, 0)
@@ -504,7 +506,8 @@ def _softmax_deprecated(
     where: ArrayLike | None = None,
     initial: ArrayLike | None = None) -> Array:
   x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True)
-  unnormalized = jnp.exp(x - lax.stop_gradient(x_max))
+  x_safe = x if where is None else jnp.where(where, x, initial)
+  unnormalized = jnp.exp(x_safe - lax.stop_gradient(x_max))
   result = unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
   if where is not None:
     result = jnp.where(where, result, 0)


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -172,6 +172,16 @@ def testSoftmaxWhereMask(self, fn):
                                 jnp.array([0, 2, 3]))
       jtu.check_grads(g_fun, (x,), order=2)
 
+  @parameterized.parameters([nn.softmax, nn.log_softmax])
+  def testSoftmaxWhereGrad(self, fn):
+    # regression test for https://github.com/google/jax/issues/19490
+    x = jnp.array([36., 10000.])
+    mask = x < 1000
+
+    f = lambda x, mask: fn(x, where=mask, initial=x.min())[0]
+
+    self.assertAllClose(jax.grad(f)(x, mask), jnp.zeros_like(x))
+
   def testSoftmaxGrad(self):
     x = jnp.array([5.5, 1.3, -4.2, 0.9])
     jtu.check_grads(nn.softmax, (x,), order=2, atol=5e-3)