Repository: google/jax (python)

Issue:
Incorrect Jacobian of `jax.scipy.special.logsumexp` when `b` contains a `0`
### Description

The gradient of `jax.scipy.special.logsumexp` along `b` is wrong when `b` contains a `0`. In the following example, the gradient of $f(x_1, x_2) = \log(x_1 + x_2)$ should be:
```math
\nabla_x f(x) = \left( \frac{1}{x_1 + x_2}, \frac{1}{x_1 + x_2} \right)
```
but JAX outputs that the gradient at `(1,0)` is `(1,0)` when it should be `(1,1)`:

```python
import jax
import jax.numpy as jnp

fun = lambda x: jax.scipy.special.logsumexp(jnp.zeros(2), axis=0, b=x)
a = jnp.array([1, 0], dtype=float)

jax.jacfwd(fun)(a)
```
output:
```python
Array([1., 0.], dtype=float32)
```


There is only a problem for an exact `0`: the gradient taken at `(1,1e-30)` outputs the correct value of `(1,1)`:
```python
a = jnp.array([1, 1e-30], dtype=float)
jax.jacfwd(fun)(a)
```
output:
```python
Array([1., 1.], dtype=float32)
```


Furthermore, if the `logsumexp` is implemented manually, the gradient at `(1, 0)` is also correct:
```python
fun_2 = lambda x: jnp.log(jnp.sum(x * jnp.exp(jnp.zeros(2))))
a = jnp.array([1, 0], dtype=float)
jax.jacfwd(fun_2)(a)
```
output:
```python
Array([1., 1.], dtype=float32)
```

cc @PhilipVinc

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.30
jaxlib: 0.4.30
numpy:  1.24.4
python: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:37:07) [Clang 15.0.7 ]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='mba-10836208', release='23.4.0', version='Darwin Kernel Version 23.4.0: Fri Mar 15 00:19:22 PDT 2024; root:xnu-10063.101.17~1/RELEASE_ARM64_T8112', machine='arm64')
```

Pull request description:
Rolling forward #22836
Rolling forward #22836

This version, proposed by @dfm, does not have a custom JVP for the whole 
logsumexp and instead fixes #22398 directly.

Reverts e416c6675acfd82866a6e83e8c221640c4d02f29


Patch:
diff --git a/jax/_src/ops/special.py b/jax/_src/ops/special.py
index 59ad594ef..45b26b0de 100644
--- a/jax/_src/ops/special.py
+++ b/jax/_src/ops/special.py
@@ -14,12 +14,12 @@
 
 from __future__ import annotations
 
-from typing import overload, Literal
+from typing import Literal, overload
 
 import jax
 from jax import lax
 from jax import numpy as jnp
-from jax._src.numpy.reductions import _reduction_dims, Axis
+from jax._src.numpy.reductions import Axis, _reduction_dims
 from jax._src.numpy.util import promote_args_inexact
 from jax._src.typing import Array, ArrayLike
 import numpy as np
@@ -40,6 +40,7 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
 def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
               keepdims: bool = False, return_sign: bool = False, where: ArrayLike | None = None) -> Array | tuple[Array, Array]: ...
 
+
 def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
               keepdims: bool = False, return_sign: bool = False, where: ArrayLike | None = None) -> Array | tuple[Array, Array]:
   r"""Log-sum-exp reduction.
@@ -71,18 +72,22 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
   """
   if b is not None:
     a_arr, b_arr = promote_args_inexact("logsumexp", a, b)
-    a_arr = jnp.where(b_arr != 0, a_arr, -jnp.inf)
+    a_masked = jnp.where(b_arr != 0, a_arr, -jnp.inf)
   else:
     a_arr, = promote_args_inexact("logsumexp", a)
     b_arr = a_arr  # for type checking
+    a_masked = a_arr
   pos_dims, dims = _reduction_dims(a_arr, axis)
-  amax = jnp.max(a_arr.real, axis=dims, keepdims=keepdims, where=where, initial=-jnp.inf)
+  amax = jnp.max(
+      a_masked.real, axis=dims, keepdims=keepdims, where=where, initial=-jnp.inf
+  )
   amax = lax.stop_gradient(lax.select(jnp.isfinite(amax), amax, lax.full_like(amax, 0)))
   amax_with_dims = amax if keepdims else lax.expand_dims(amax, pos_dims)
 
-  exp_a = lax.exp(lax.sub(a_arr, amax_with_dims.astype(a_arr.dtype)))
-  if b is not None:
-    exp_a = lax.mul(exp_a, b_arr)
+  if b is None:
+    exp_a = lax.exp(lax.sub(a_arr, amax_with_dims.astype(a_arr.dtype)))
+  else:
+    exp_a = _stable_mulexp(a_arr - amax_with_dims.astype(a_arr.dtype), b_arr)
   sumexp = exp_a.sum(axis=dims, keepdims=keepdims, where=where)
   sign = lax.sign(sumexp)
   if return_sign or not np.issubdtype(a_arr.dtype, np.complexfloating):
@@ -95,3 +100,20 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
     with jax.debug_nans(False):
       out = jnp.where(sign < 0, jnp.array(np.nan, dtype=out.dtype), out)
   return out
+
+
+@jax.custom_jvp
+def _stable_mulexp(a_scaled: Array, b: Array) -> Array:
+  # This helper ensures that the output of logsumexp depends on b for b == 0.
+  # See https://github.com/google/jax/issues/22398.
+  a_scaled = jnp.where(b != 0, a_scaled, -jnp.inf)
+  return lax.mul(lax.exp(a_scaled), b)
+
+
+@_stable_mulexp.defjvp
+def _stable_mulexp_jvp(primals, tangents):
+  a_scaled, b = primals
+  da, db = tangents
+  out = _stable_mulexp(a_scaled, b)
+  dout = _stable_mulexp(a_scaled, db) + da * out
+  return out, dout


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index 1ed410cba..66d84c427 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -20,9 +20,9 @@ import unittest
 from absl.testing import absltest
 
 import numpy as np
+import scipy.cluster as osp_cluster
 import scipy.integrate
 import scipy.special as osp_special
-import scipy.cluster as osp_cluster
 
 import jax
 import jax.dtypes
@@ -202,6 +202,14 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     y_actual = lsp_special.logsumexp(x, where=mask)
     self.assertAllClose(y_expected, y_actual, check_dtypes=False)
 
+  def testLogSumExpZerosJac(self):
+    # Regression test for https://github.com/google/jax/issues/22398
+    fun = lambda b: lsp_special.logsumexp(jnp.zeros(2), axis=0, b=b)
+    np.testing.assert_array_equal(
+        jax.jacfwd(fun)(jnp.array([1.0, 0.0])),
+        jnp.ones(2),
+    )
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=float_dtypes,