Repository: scipy/scipy
Pull request: ENH: `special.logsumexp`: improve precision when one element is much bigger than the rest
#### Reference issue
Closes gh-18295
Supersedes gh-18424
~~May be used to address gh-19521/gh-19549 and make `log_softmax` array API compatible.~~ (The idea of using `log1p` is essentially the same, but I think `softmax` needs its own implementation.)

#### What does this implement/fix?
gh-18295 reported that `logsumexp` can lose precision when one element is much bigger than the rest, especially when the exponential of it is close to 1. This improves the precision as described in the issue and [linked paper](https://academic.oup.com/imajna/article/41/4/2311/5893596?login=true).

#### Additional information
gh-18424 was out of date after converting `logsumexp` to the array API. For instance, `xp.max` does not work on the real component of complex arrays, so conversion of some parts would not be trivial. Also, there were some unresolved comments about the complexity, so I chose to start from scratch.

Also, `logsumexp` was getting quite complicated as it was. I found it challenging to work within the existing structure, so I refactored to simplify (26bb631f4c015a1ddea1c77f1c4f5a9ad2ebc28d) before getting started with the upgrade. 

I'll add a review that documents the math inline with the code.

Problem statement:
BUG: special: Loss of precision in `logsumexp`
### Describe your issue.

`logsumexp` loses precision in cases where it shouldn't, and in extreme cases it can return 0 when it should return a nonzero float.

### Reproducing Code Example

```python
In [2]: from scipy.special import logsumexp

In [3]: logsumexp([0, -40.0])
Out[3]: 0.0

In [4]: np.logaddexp(0, -40)  # This is the expected result of `logsumexp([0, -40.0])`.
Out[4]: 4.248354255291589e-18
```

<details>
<summary>

### SciPy/NumPy/Python version and system information

</summary>

```shell
{
  "Compilers": {
    "c": {
      "name": "gcc",
      "linker": "ld.bfd",
      "version": "11.3.0",
      "commands": "cc"
    },
    "cython": {
      "name": "cython",
      "linker": "cython",
      "version": "0.29.33",
      "commands": "cython"
    },
    "c++": {
      "name": "gcc",
      "linker": "ld.bfd",
      "version": "11.3.0",
      "commands": "c++"
    },
    "fortran": {
      "name": "gcc",
      "linker": "ld.bfd",
      "version": "11.3.0",
      "commands": "gfortran"
    },
    "pythran": {
      "version": "0.12.1",
      "include directory": "../../../../../py3.10.8/lib/python3.10/site-packages/pythran"
    }
  },
  "Machine Information": {
    "host": {
      "cpu": "x86_64",
      "family": "x86_64",
      "endian": "little",
      "system": "linux"
    },
    "build": {
      "cpu": "x86_64",
      "family": "x86_64",
      "endian": "little",
      "system": "linux"
    },
    "cross-compiled": false
  },
  "Build Dependencies": {
    "blas": {
      "name": "openblas",
      "found": true,
      "version": "0.3.20",
      "detection method": "pkgconfig",
      "include directory": "/usr/include/x86_64-linux-gnu/openblas-pthread/",
      "lib directory": "/usr/lib/x86_64-linux-gnu/openblas-pthread/",
      "openblas configuration": "USE_64BITINT= DYNAMIC_ARCH=1 DYNAMIC_OLDER=1 NO_CBLAS= NO_LAPACK= NO_LAPACKE=1 NO_AFFINITY=1 USE_OPENMP=0 generic MAX_THREADS=64",
      "pc file directory": "/usr/lib/x86_64-linux-gnu/pkgconfig"
    },
    "lapack": {
      "name": "openblas",
      "found": true,
      "version": "0.3.20",
      "detection method": "pkgconfig",
      "include directory": "/usr/include/x86_64-linux-gnu/openblas-pthread/",
      "lib directory": "/usr/lib/x86_64-linux-gnu/openblas-pthread/",
      "openblas configuration": "USE_64BITINT= DYNAMIC_ARCH=1 DYNAMIC_OLDER=1 NO_CBLAS= NO_LAPACK= NO_LAPACKE=1 NO_AFFINITY=1 USE_OPENMP=0 generic MAX_THREADS=64",
      "pc file directory": "/usr/lib/x86_64-linux-gnu/pkgconfig"
    }
  },
  "Python Information": {
    "path": "/home/warren/py3.10.8/bin/python3",
    "version": "3.10"
  }
}
```

</details>

Hints from issue discussion:
The following script contains a couple alternative implementations for `logsumexp(a)` that avoid the loss of precision.  These are for a 1-d input with no weight parameter; actually implementing `scipy.special.logsumexp` with one of these will require some more work.

```
import numpy as np
from mpmath import mp


def mp_naive_logsumexp(a):
    with mp.extradps(10):
        return mp.log(mp.fsum([mp.exp(ai) for ai in a]))


def np_logaddexp_reduce_sort(a):
    return np.logaddexp.reduce(np.sort(a)[::-1])


def logsumexp_v2(a):
    a = np.asarray(a)
    a_max = a.max()
    b = a[a != a_max] - a_max
    m = len(a) - len(b)
    return np.log1p(np.exp(b).sum()/m) + np.log(m) + a_max
```
`mp_naive_logsumexp` provides a reference implementation based on `mpmath`.  Be sure to set `mp.dps` sufficiently high to avoid loss of precision in that implementation.

The simpler implementation is just `np.logaddexp.reduce(np.sort(a)[::-1])`. E.g.

```
In [95]: a = [0, -40]

In [96]: np_logaddexp_reduce_sort(a)
Out[96]: 4.248354255291589e-18
```

The array is sorted in descending order in that implementation to improve the result for an input such as `a = [-600, -600, -600, -600, -600, -600, -600, 0]`.  If the array is not sorted before calling `np.logaddexp.reduce`, the result is `1.8552775871029815e-260`, which has a relative error of about 2e-14.  If the array is sorted, the result is accurate to machine precision:

```
In [171]: a = [-600, -600, -600, -600, -600, -600, -600, 0]

In [172]: mp.dps = 400

In [173]: float(mp_naive_logsumexp(a))
Out[173]: 1.8552775871030176e-260

In [174]: np_logaddexp_reduce_sort(a)
Out[174]: 1.8552775871030176e-260
```
(Such high precision is not guaranteed in all cases, of course.)

The other implementation, `logsumexp_v2`, is a refinement of the current implementation in `logsumexp`.  Let $\hat{a}$ be the maximum value in `a`.  Then

$$\begin{align}
  \log\left(\sum_i e^{a_i}\right)
      &= \log\left(\sum_{i}e^{\hat{a}}e^{a_i - \hat{a}}\right) \\
      &= \log\left(e^{\hat{a}} \sum_{i} e^{a_i - \hat{a}}\right) \\
      &= \hat{a} + \log\left(\sum_{i}e^{a_i - \hat{a}}\right) \\
      &= \hat{a} + \log\left(m + \sum_{i \in S} e^{a_i - \hat{a}}\right) \\
      &= \hat{a} + \log(m) + \log\left(1 + \frac{1}{m}\sum_{i \in S} e^{a_i - \hat{a}}\right)
\end{align}$$

where $m$ is the number of elements in $a$ that equal $\hat{a}$ and $S$ is the set of indices of $a$ where $a_i \ne \hat{a}$.  The first three lines in the above equations are the standard "trick" of factoring out the maximum of `a`; that is what is implemented currently in `logsumexp`.  In the remaining equations, the values in the sum that are known to be 1 are pulled out of the sum, allowing `log1p` to be used in the implementation.

The implementation in `logsumexp_v2` handles the examples shown above correctly:
```
In [498]: a = [0, -40]

In [499]: logsumexp_v2(a)
Out[499]: 4.248354255291589e-18

In [500]: a = [-600, -600, -600, -600, -600, -600, -600, 0]

In [501]: logsumexp_v2(a)
Out[501]: 1.8552775871030176e-260
```

`logsumexp_v2` appears to work as well as `np_logaddexp_reduce_sort`, and possibly even slightly better.  In my experiments, both generally get within a few ULPs of the reference value, but `logsumexp_v2` might do slightly better more often.  But that might be a bias created by the various specific examples I've experimented with.

I probably won't get around to creating a pull request based on one of these implementations to fix `logsumexp` any time soon, so anyone is welcome to do so.  You'll have to deal with the `axis` and `b` parameters appropriately.

Commit message: TST: integrate._tanhsinh: update test tolerance

Patch to review:
diff --git a/scipy/integrate/tests/test_tanhsinh.py b/scipy/integrate/tests/test_tanhsinh.py
index ee990bd0e787..bfeff6d21ddd 100644
--- a/scipy/integrate/tests/test_tanhsinh.py
+++ b/scipy/integrate/tests/test_tanhsinh.py
@@ -844,7 +844,7 @@ def test_basic(self, f_number):
             logres = nsum(lambda *args: np.log(f(*args)),
                            f.a, f.b, log=True, args=f.args)
         assert_allclose(np.exp(logres.sum), res.sum)
-        assert_allclose(np.exp(logres.error), res.error)
+        assert_allclose(np.exp(logres.error), res.error, atol=1e-15)
         assert_equal(logres.status, 0)
         assert_equal(logres.success, True)
 
diff --git a/scipy/special/_logsumexp.py b/scipy/special/_logsumexp.py
index d8efad6d5bde..bd4d7f82de63 100644
--- a/scipy/special/_logsumexp.py
+++ b/scipy/special/_logsumexp.py
@@ -4,6 +4,7 @@
     array_namespace,
     xp_size,
     xp_broadcast_promote,
+    xp_atleast_nd,
 )
 
 __all__ = ["logsumexp", "softmax", "log_softmax"]
@@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
-        # because `xp.max` doesn't have `initial` argument...
-        shape = np.asarray(a.shape)  # NumPy is concise for scalar or tuple `axis`
+        shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
+        axis = tuple(range(a.ndim)) if axis is None else axis
         shape[axis] = 1
-        a_max = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        sgn = xp.sign(out)
+    else:
+        with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
+            out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
+
+    # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy
+    out = xp.squeeze(out, axis=axis) if not keepdims else out
+    sgn = xp.squeeze(sgn, axis=axis) if (sgn is not None and not keepdims) else sgn
+    out = out[()] if out.ndim == 0 else out
+    sgn = sgn[()] if (sgn is not None and sgn.ndim == 0) else sgn
+
+    return (out, sgn) if return_sign else out
+
+
+def _elements_and_indices_with_max_real(a, axis=-1, xp=None):
+    # We need *an* element with the maximum real part.
+    # If the array is complex, we need to choose *one* of the complex values,
+    # not all values with the same real part.
+    # We could use `argmax` on real component, but array API doesn't yet have
+    # `take_along_axis`, and even if it did, we would have problems with axis tuples.
+    xp = array_namespace(a) if xp is None else xp
+
+    if xp.isdtype(a.dtype, "complex floating"):
+        # select all elements with max real part.
+        max = xp.max(xp.real(a), axis=axis, keepdims=True)
+        mask = xp.real(a) == max
+
+        # Of those, choose one arbitrarily. This is a reasonably
+        # simple, array-API compatible way of doing so that doesn't
+        # have a problem with `axis` being a tuple or None.
+        i = xp.reshape(xp.arange(xp_size(a)), a.shape)
+        i[~mask] = -1
+        max_i = xp.max(i, axis=axis, keepdims=True)
+        mask = i == max_i
+        max = xp.reshape(a[mask], max_i.shape)
     else:
-        real = xp.real(a) if xp.isdtype(a.dtype, "complex floating") else a
-        a_max = xp.max(real, axis=axis, keepdims=True)
+        max = xp.max(a, axis=axis, keepdims=True)
+        mask = a == max
 
-    if a_max.ndim > 0:
-        a_max[~xp.isfinite(a_max)] = 0
-    elif not xp.isfinite(a_max):
-        a_max = 0
+    return xp.asarray(max), xp.asarray(mask)
 
-    if b is not None:
-        b = xp.asarray(b)
-        tmp = b * xp.exp(a - a_max)
-    else:
-        tmp = xp.exp(a - a_max)
 
-    # suppress warnings about log of zero
-    with np.errstate(divide='ignore'):
-        s = xp.sum(tmp, axis=axis, keepdims=keepdims, dtype=tmp.dtype)
-        if return_sign:
-            # For complex, use the numpy>=2.0 convention for sign.
-            if xp.isdtype(s.dtype, "complex floating"):
-                sgn = s / xp.where(s == 0, xp.asarray(1, dtype=s.dtype), xp.abs(s))
-            else:
-                sgn = xp.sign(s)
-            s = xp.abs(s)
-        out = xp.log(s)
-
-    if not keepdims:
-        a_max = xp.squeeze(a_max, axis=axis)
-    out += a_max
+def _sign(x, xp):
+    return x / xp.where(x == 0, xp.asarray(1, dtype=x.dtype), xp.abs(x))
+
 
+def _logsumexp(a, b, axis, return_sign, xp):
+
+    # This has been around for about a decade, so let's consider it a feature:
+    # Even if element of `a` is infinite or NaN, it adds nothing to the sum if
+    # the corresponding weight is zero.
+    if b is not None:
+        a[b == 0] = -xp.inf
+
+    # Find element with maximum real part, since this is what affects the magnitude
+    # of the exponential. Possible enhancement: include log of `b` magnitude in `a`.
+    a_max, i_max = _elements_and_indices_with_max_real(a, axis=axis, xp=xp)
+
+    # for precision, these terms are separated out of the main sum.
+    a[i_max] = -xp.inf
+    i_max_dt = xp.astype(i_max, a.dtype)
+    # This is an inefficient way of getting `m` because it is the sum of a sparse
+    # array; however, this is the simplest way I can think of to get the right shape.
+    m = (xp.sum(i_max_dt, axis=axis, keepdims=True, dtype=a.dtype) if b is None
+         else xp.sum(b * i_max_dt, axis=axis, keepdims=True, dtype=a.dtype))
+
+    # Arithmetic between infinities will introduce NaNs.
+    # `+ a_max` at the end naturally corrects for removing them here.
+    shift = xp.where(xp.isfinite(a_max), a_max, xp.asarray(0, dtype=a_max.dtype))
+
+    # Shift, exponentiate, scale, and sum
+    exp = b * xp.exp(a - shift) if b is not None else xp.exp(a - shift)
+    s = xp.sum(exp, axis=axis, keepdims=True, dtype=exp.dtype)
+    s = xp.where(s == 0, s, s/m)
+
+    # Separate sign/magnitude information
+    sgn = None
     if return_sign:
-        return out, sgn
-    else:
-        return out
+        # Use the numpy>=2.0 convention for sign.
+        # When all array libraries agree, this can become sng = xp.sign(s).
+        sgn = _sign(s + 1, xp=xp) * _sign(m, xp=xp)
+
+        if xp.isdtype(s.dtype, "real floating"):
+            # The log functions need positive arguments
+            s = xp.where(s < -1, -s - 2, s)
+            m = xp.abs(m)
+        else:
+            # `a_max` can have a sign component for complex input
+            j = xp.asarray(1j, dtype=a_max.dtype)
+            sgn = sgn * xp.exp(xp.imag(a_max) * j)
+
+    # Take log and undo shift
+    out = xp.log1p(s) + xp.log(m) + a_max
+
+    # Sigh... imagine a world in which we could just take the real part of a real array
+    out = (xp.real(out) if return_sign and xp.isdtype(s.dtype, "complex floating")
+           else out)
+
+    return out, sgn
 
 
 def softmax(x, axis=None):
diff --git a/scipy/special/tests/test_logsumexp.py b/scipy/special/tests/test_logsumexp.py
index 138959a5bc47..3b3306b730c5 100644
--- a/scipy/special/tests/test_logsumexp.py
+++ b/scipy/special/tests/test_logsumexp.py
@@ -195,6 +195,12 @@ def test_dtypes_ab(self, dtype_a, dtype_b, xp):
         desired = xp.asarray(math.log(math.exp(2) - math.exp(1)), dtype=desired_dtype)
         xp_assert_close(logsumexp(a, b=b), desired)
 
+    def test_gh18295(self, xp):
+        a = xp.asarray([0.0, -40.0])
+        res = logsumexp(a)
+        ref = xp.logaddexp(a[0], a[1])
+        xp_assert_close(res, ref)
+
 
 class TestSoftmax:
     def test_softmax_fixtures(self):


Code review comments:
diff_hunk: @@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
-        # because `xp.max` doesn't have `initial` argument...
-        shape = np.asarray(a.shape)  # NumPy is concise for scalar or tuple `axis`
+        shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
+        axis = tuple(range(a.ndim)) if axis is None else axis
         shape[axis] = 1
-        a_max = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        sgn = xp.sign(out)
+    else:
+        with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
+            out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
+
+    # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy
+    out = xp.squeeze(out, axis=axis) if not keepdims else out
+    sgn = xp.squeeze(sgn, axis=axis) if (sgn is not None and not keepdims) else sgn
+    out = out[()] if out.ndim == 0 else out
+    sgn = sgn[()] if (sgn is not None and sgn.ndim == 0) else sgn
+
+    return (out, sgn) if return_sign else out
+
+
+def _elements_and_indices_with_max_real(a, axis=-1, xp=None):
+    # We need *an* element with the maximum real part.
+    # If the array is complex, we need to choose *one* of the complex values,
+    # not all values with the same real part.
+    # We could use `argmax` on real component, but array API doesn't yet have
+    # `take_along_axis`, and even if it did, we would have problems with axis tuples.
+    xp = array_namespace(a) if xp is None else xp
+
+    if xp.isdtype(a.dtype, "complex floating"):
+        # select all elements with max real part.
+        max = xp.max(xp.real(a), axis=axis, keepdims=True)
+        mask = xp.real(a) == max
+
+        # Of those, choose one arbitrarily. This is a reasonably
+        # simple, array-API compatible way of doing so that doesn't
+        # have a problem with `axis` being a tuple or None.
+        i = xp.reshape(xp.arange(xp_size(a)), a.shape)
+        i[~mask] = -1
+        max_i = xp.max(i, axis=axis, keepdims=True)
+        mask = i == max_i
+        max = xp.reshape(a[mask], max_i.shape)
path: scipy/special/_logsumexp.py
text: @user1:
Can't say this is very intuitive to read :/

@author:
Sorry. Sometimes working around the Array API standard is tricky. Feel free to suggest another implementation that behaves like `np.max` that is Array API Standard compatible.

diff_hunk: @@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
-        # because `xp.max` doesn't have `initial` argument...
-        shape = np.asarray(a.shape)  # NumPy is concise for scalar or tuple `axis`
+        shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
+        axis = tuple(range(a.ndim)) if axis is None else axis
         shape[axis] = 1
-        a_max = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        sgn = xp.sign(out)
+    else:
+        with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
+            out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
+
+    # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy
+    out = xp.squeeze(out, axis=axis) if not keepdims else out
+    sgn = xp.squeeze(sgn, axis=axis) if (sgn is not None and not keepdims) else sgn
+    out = out[()] if out.ndim == 0 else out
+    sgn = sgn[()] if (sgn is not None and sgn.ndim == 0) else sgn
+
+    return (out, sgn) if return_sign else out
+
+
+def _elements_and_indices_with_max_real(a, axis=-1, xp=None):
+    # We need *an* element with the maximum real part.
+    # If the array is complex, we need to choose *one* of the complex values,
+    # not all values with the same real part.
+    # We could use `argmax` on real component, but array API doesn't yet have
+    # `take_along_axis`, and even if it did, we would have problems with axis tuples.
+    xp = array_namespace(a) if xp is None else xp
+
+    if xp.isdtype(a.dtype, "complex floating"):
+        # select all elements with max real part.
+        max = xp.max(xp.real(a), axis=axis, keepdims=True)
+        mask = xp.real(a) == max
+
+        # Of those, choose one arbitrarily. This is a reasonably
+        # simple, array-API compatible way of doing so that doesn't
+        # have a problem with `axis` being a tuple or None.
+        i = xp.reshape(xp.arange(xp_size(a)), a.shape)
+        i[~mask] = -1
+        max_i = xp.max(i, axis=axis, keepdims=True)
+        mask = i == max_i
+        max = xp.reshape(a[mask], max_i.shape)
     else:
-        real = xp.real(a) if xp.isdtype(a.dtype, "complex floating") else a
-        a_max = xp.max(real, axis=axis, keepdims=True)
+        max = xp.max(a, axis=axis, keepdims=True)
+        mask = a == max
 
-    if a_max.ndim > 0:
-        a_max[~xp.isfinite(a_max)] = 0
-    elif not xp.isfinite(a_max):
-        a_max = 0
+    return xp.asarray(max), xp.asarray(mask)
 
-    if b is not None:
-        b = xp.asarray(b)
-        tmp = b * xp.exp(a - a_max)
-    else:
-        tmp = xp.exp(a - a_max)
 
-    # suppress warnings about log of zero
-    with np.errstate(divide='ignore'):
-        s = xp.sum(tmp, axis=axis, keepdims=keepdims, dtype=tmp.dtype)
-        if return_sign:
-            # For complex, use the numpy>=2.0 convention for sign.
-            if xp.isdtype(s.dtype, "complex floating"):
-                sgn = s / xp.where(s == 0, xp.asarray(1, dtype=s.dtype), xp.abs(s))
-            else:
-                sgn = xp.sign(s)
-            s = xp.abs(s)
-        out = xp.log(s)
-
-    if not keepdims:
-        a_max = xp.squeeze(a_max, axis=axis)
-    out += a_max
+def _sign(x, xp):
+    return x / xp.where(x == 0, xp.asarray(1, dtype=x.dtype), xp.abs(x))
+
 
+def _logsumexp(a, b, axis, return_sign, xp):
+
+    # This has been around for about a decade, so let's consider it a feature:
+    # Even if element of `a` is infinite or NaN, it adds nothing to the sum if
+    # the corresponding weight is zero.
+    if b is not None:
+        a[b == 0] = -xp.inf
+
+    # Find element with maximum real part, since this is what affects the magnitude
+    # of the exponential. Possible enhancement: include log of `b` magnitude in `a`.
+    a_max, i_max = _elements_and_indices_with_max_real(a, axis=axis, xp=xp)
+
+    # for precision, these terms are separated out of the main sum.
+    a[i_max] = -xp.inf
+    i_max_dt = xp.astype(i_max, a.dtype)
+    # This is an inefficient way of getting `m` because it is the sum of a sparse
+    # array; however, this is the simplest way I can think of to get the right shape.
+    m = (xp.sum(i_max_dt, axis=axis, keepdims=True, dtype=a.dtype) if b is None
+         else xp.sum(b * i_max_dt, axis=axis, keepdims=True, dtype=a.dtype))
+
+    # Arithmetic between infinities will introduce NaNs.
+    # `+ a_max` at the end naturally corrects for removing them here.
+    shift = xp.where(xp.isfinite(a_max), a_max, xp.asarray(0, dtype=a_max.dtype))
+
+    # Shift, exponentiate, scale, and sum
+    exp = b * xp.exp(a - shift) if b is not None else xp.exp(a - shift)
+    s = xp.sum(exp, axis=axis, keepdims=True, dtype=exp.dtype)
+    s = xp.where(s == 0, s, s/m)
+
+    # Separate sign/magnitude information
+    sgn = None
     if return_sign:
-        return out, sgn
-    else:
-        return out
+        # Use the numpy>=2.0 convention for sign.
+        # When all array libraries agree, this can become sng = xp.sign(s).
+        sgn = _sign(s + 1, xp=xp) * _sign(m, xp=xp)
+
+        if xp.isdtype(s.dtype, "real floating"):
+            # The log functions need positive arguments
+            s = xp.where(s < -1, -s - 2, s)
+            m = xp.abs(m)
+        else:
+            # `a_max` can have a sign component for complex input
+            j = xp.asarray(1j, dtype=a_max.dtype)
+            sgn = sgn * xp.exp(xp.imag(a_max) * j)
+
+    # Take log and undo shift
+    out = xp.log1p(s) + xp.log(m) + a_max
+
+    # Sigh... imagine a world in which we could just take the real part of a real array
path: scipy/special/_logsumexp.py
text: @user1:
:shrug: Works OK with NumPy?

@author:
data-apis/array-api#824

@user2:
Remember that we have encountered this before, and added `xp_real` in `_lib`.

@author:
Well you can see that I  remembered encountering it before, but I did forget that I added `xp_real`. I'll use it.

diff_hunk: @@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
path: scipy/special/_logsumexp.py
text: @author:
Make size-0 arrays a special case. It possible that they can be made to work with the existing code again. I'd be happy to review that as a simple follow-up. But conceptually, it's simpler to treat the special case here so we don't need to complicate the algorithmic code.

diff_hunk: @@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
-        # because `xp.max` doesn't have `initial` argument...
-        shape = np.asarray(a.shape)  # NumPy is concise for scalar or tuple `axis`
+        shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
+        axis = tuple(range(a.ndim)) if axis is None else axis
         shape[axis] = 1
-        a_max = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        sgn = xp.sign(out)
+    else:
+        with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
+            out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
path: scipy/special/_logsumexp.py
text: @user1:
I'd have swapped these lines so the most common behavior occurs first, and the exceptional case (`xp_size(a) == 0`) comes later. It took me a bit to figure out what the function is doing!

@author:
I can swap these for you if you prefer.

diff_hunk: @@ -100,53 +101,113 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """
     xp = array_namespace(a, b)
     a, b = xp_broadcast_promote(a, b, ensure_writeable=True, force_floating=True, xp=xp)
-    axis = tuple(range(a.ndim)) if axis is None else axis
+    a = xp_atleast_nd(a, ndim=1)
+    b = xp_atleast_nd(b, ndim=1) if b is not None else b
 
-    if b is not None:
-        a[b == 0] = -xp.inf
-
-    # Scale by real part for complex inputs, because this affects
-    # the magnitude of the exponential.
     if xp_size(a) == 0:
-        # because `xp.max` doesn't have `initial` argument...
-        shape = np.asarray(a.shape)  # NumPy is concise for scalar or tuple `axis`
+        shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
+        axis = tuple(range(a.ndim)) if axis is None else axis
         shape[axis] = 1
-        a_max = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        out = xp.full(tuple(shape), -xp.inf, dtype=a.dtype)
+        sgn = xp.sign(out)
+    else:
+        with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
+            out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
+
+    # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy
+    out = xp.squeeze(out, axis=axis) if not keepdims else out
+    sgn = xp.squeeze(sgn, axis=axis) if (sgn is not None and not keepdims) else sgn
+    out = out[()] if out.ndim == 0 else out
+    sgn = sgn[()] if (sgn is not None and sgn.ndim == 0) else sgn
+
+    return (out, sgn) if return_sign else out
+
+
+def _elements_and_indices_with_max_real(a, axis=-1, xp=None):
+    # We need *an* element with the maximum real part.
path: scipy/special/_logsumexp.py
text: @user1:
This docstring isn't immediately intuitive to me. What is "*an* element with the maximum real part"? Does that mean, any element that has a real part equal to the maximum found in the array?

@author:
We need a function that behaves like behaves like np.max with complex input that is Array API Standard compatible.
```python3
np.max([1 + 5j, 2 + 1j, 2 + 0j])
# np.complex128(2+1j), but np.complex128(2+ 0j) would be OK, too. 
# I don't care about the imaginary part
```
It needs to work for n-d arrays and axis tuples/None.

@user1:
Maybe we need a shim that offers an Array API equivalent of all these functions, so we can just write in NumPy parlance? Such a thing probably already exists. I just imagine it's going to be rather confusing to re-write a bunch of code in ways that handle all the Array API differences, and do that repeatedly for the various occurrences.

@author:
Yes, this idea has been considered. When I was first introduced to `array_api_compat`, this is exactly what I expected it to do. From an adoption standpoint, it would be much easier if the library just did translations from NumPy/CuPy syntax to calls in other backends, since SciPy (and many other libraries) is written for NumPy. I was very surprised that this *wasn't* what it did, but I have come to terms with it. I understand that the vision is for a cleaner, universal syntax rather than layers upon layers of translations from an API with 25 years of accumulated patches.

In any case, for a while now we have been converting code to use the standard directly, making a few of our own convenience functions where particularly important. See e.g. https://github.com/data-apis/array-api/issues/837#issuecomment-2361816873 for a related idea that has been mulled over for a long time.

So whether to have a NumPy to array API shim library or not is a bit beyond this PR; perhaps bring it up in gh-18286 if it hasn't been discussed there explicitly.

(As for whether I think we should use a shim library, I am conflicted at this point. Clearly it would be easier to "translate" if we used a shim. On these other hand, these translations are forcing us to review and modernize some old code. When we review, we sometimes find bugs and inefficiencies that we can fix. In some places, the standard forces us to jump through hoops that seem unnecessary, but with a few tweaks to the standard - which we suggest as we find these hoops repeatedly - I think we could end up with a much cleaner and more maintainable library.)

---

As for this function, I would be happy for it to be rewritten. Since it was just an annoyance that it didn't already exist, I did it in the first way that came to mind that works given the constraints. 

The cleaner way would probably be to standardize `axis` (ravel axis tuples into one axis), then use `argmax` / `take_along_axis`, then reshape. But this is tricky, too, and it relies on `take_along_axis` which was only added to the standard a few days ago, so it hasn't made its way into `array_api_compat`.

@user2:
> ```python
> np.max([1 + 5j, 2 + 1j, 2 + 0j])
> # np.complex128(2+1j), but np.complex128(2+ 0j) would be OK, too. 
> # I don't care about the imaginary part
> ```

would you be okay with spelling such a function like 

```python
max(xp.asarray([1 + 5j, 2 + 1j, 2 + 0j]), complex_ok=True, xp=xp)
```

? If so, since this could be implemented in terms of standard operations, it would be in-scope for https://github.com/data-apis/array-api-extra.

---

EDIT: 

> (As for whether I think we should use a shim library, I am conflicted at this point. Clearly it would be easier to "translate" if we used a shim. On these other hand, these translations are forcing us to review and modernize some old code. When we review, we sometimes find bugs and inefficiencies that we can fix. In some places, the standard forces us to jump through hoops that seem unnecessary, but with a few tweaks to the standard -which we suggest as we find these hoops repeatedly - I think we could end up with a much cleaner and more maintainable library.)

I'd just like to echo these thoughts too.

@author:
> would you be okay with spelling such a function like

Sure, but I would be happy with it just being a `max` that defines a particular ordering for complex numbers; it doesn't need a `complex_ok` arg IMO if the ordering is documented. Or maybe there can be an argument to select among common rules, but with NumPy's default the default for now. Thanks for creating the library.