Repository: google/jax
Issue: Bug with masking in jax.nn.log_softmax

### Description

I recently wanted to use jax.nn.log_softmax together with a mask (where = mask).
https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.log_softmax.html

When I masked out some output probabilities from the log_sofmax I found out that masked out entries did not have zero probability.

The problem can be reproduced with the following example, where I also added my code that I used as an work around.
```
def log_softmax( logits, mask, initial = jnp.finfo(jnp.float32).min):
    masked_logits = jnp.where(mask, logits, initial * jnp.ones_like(logits))
    max_logits = jax.lax.stop_gradient(jnp.max(masked_logits, axis=-1))
    shifted = masked_logits - max_logits[:, jnp.newaxis]
    shifted_logsumexp = jnp.log( jnp.sum(jnp.exp(shifted), axis = -1, keepdims=True))
    res = shifted - shifted_logsumexp
    return res

def log_softmax_bug():
    initial = jnp.finfo(jnp.float32).min
    key = jax.random.PRNGKey(0)
    key, subkey = jax.random.split(key)
    shape = (3,10)
    logits = jax.random.normal(subkey, shape= shape)
    key, subkey = jax.random.split(key)
    mask = jax.random.randint(key, (3,10), 0, 2)

    masked_log_probs = jax.nn.log_softmax(logits, where = mask, initial = initial)


    my_masked_log_probs = log_softmax(logits, mask, initial = initial)

    print("jax nn probs", np.exp(masked_log_probs))
    print("sum probs",np.sum(np.exp(masked_log_probs), axis=-1))
    print("my probs", np.exp(my_masked_log_probs))
    print("sum probs",np.sum(np.exp(my_masked_log_probs), axis = -1))

```

The output of the code is down below, where it is clear that the probabilities of the probabilities that are produced by jax.nn.log_softmax do not sum to one.

jax nn probs [[0.65082145 0.12884209 0.09259639 0.0914019  0.01950807 0.11102043
  0.03726786 0.07620961 0.06663671 0.10840429]
 [0.21390831 0.28618205 0.34040952 0.12929024 0.2339981  0.6408521
  0.18717511 0.0577253  0.15870889 0.2091582 ]
 [0.22656752 0.26329792 0.39771214 0.15017666 0.04520763 0.02898957
  0.09993661 0.41019788 0.11462383 0.11412007]]
sum probs [1.3827088 2.4574077 1.85083  ]
my probs [[0.65082145 0.         0.09259639 0.0914019  0.01950807 0.
  0.03726786 0.         0.         0.10840429]
 [0.         0.         0.34040952 0.         0.2339981  0.
  0.         0.0577253  0.15870889 0.2091582 ]
 [0.22656752 0.26329792 0.         0.         0.         0.
  0.09993661 0.41019788 0.         0.        ]]
sum probs [0.99999994 0.99999994 0.99999994]



### What jax/jaxlib version are you using?

jax v0.3.13; jaxlib v0.3.10+cuda11.cudnn82

### Which accelerator(s) are you using?

GPU

### Additional system info

Linux

### NVIDIA GPU info

_No response_

Changed files:
jax/_src/nn/functions.py, tests/nn_test.py

diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
index cbc34f3b9..3b4ca8adf 100644
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -317,7 +317,10 @@ def log_softmax(x: Array,
   shifted = x - lax.stop_gradient(x_max)
   shifted_logsumexp = jnp.log(
       jnp.sum(jnp.exp(shifted), axis, where=where, keepdims=True))
-  return shifted - shifted_logsumexp
+  result = shifted - shifted_logsumexp
+  if where is not None:
+    return jnp.where(where, result, -jnp.inf)
+  return result
 
 
 # TODO(phawkins): this jit was found to change numerics in a test. Debug this.
@@ -357,7 +360,10 @@ def _softmax(
     initial: Optional[Array] = None) -> Array:
   x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True)
   unnormalized = jnp.exp(x - x_max)
-  return unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
+  result = unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
+  if where is not None:
+    result = jnp.where(where, result, 0)
+  return result
 
 @_softmax.defjvp
 def _softmax_jvp(axis, primals, tangents):
@@ -368,7 +374,10 @@ def _softmax_jvp(axis, primals, tangents):
 def _softmax_deprecated(x, axis, where, initial):
   x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True)
   unnormalized = jnp.exp(x - lax.stop_gradient(x_max))
-  return unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
+  result = unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
+  if where is not None:
+    result = jnp.where(where, result, 0)
+  return result
 
 
 @partial(jax.jit, static_argnames=("axis",))
diff --git a/tests/nn_test.py b/tests/nn_test.py
index e558c0b6b..46554e5bc 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -133,13 +133,12 @@ class NNFunctionsTest(jtu.JaxTestCase):
   def testSoftmaxWhereMask(self, fn):
     x = jnp.array([5.5, 1.3, -4.2, 0.9])
     m = jnp.array([True, False, True, True])
-    x_filtered = jnp.take(x, jnp.array([0, 2, 3]))
 
-    out_masked = jnp.take(
-        fn(x, where=m, initial=-jnp.inf), jnp.array([0, 2, 3]))
-    out_filtered = fn(x_filtered)
+    out = fn(x, where=m, initial=-jnp.inf)
+    self.assertAllClose(out[m], fn(x[m]))
 
-    self.assertAllClose(out_masked, out_filtered)
+    probs = out if fn is nn.softmax else jnp.exp(out)
+    self.assertAllClose(probs.sum(), 1.0)
 
     # TODO(mattjj): include log_softmax in these extra tests if/when we add a
     # custom_jvp rule for it (since otherwise it doesn't pass the numerical