Merge pull request scipy#22843 from crusaderky/special_xp_capabilities

mdhaber · web-flow · commit e29a7a8eac81 · 2025-04-16T09:08:19.000-07:00
DOC: `special`: add `xp_capabilities` to logsumexp
diff --git a/scipy/special/_logsumexp.py b/scipy/special/_logsumexp.py
@@ -1,6 +1,7 @@
 import numpy as np
 from scipy._lib._array_api import (
     array_namespace,
+    xp_capabilities,
     xp_device,
     xp_size,
     xp_promote,
@@ -11,6 +12,7 @@
 __all__ = ["logsumexp", "softmax", "log_softmax"]
 
 
+@xp_capabilities()
 def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     """Compute the log of the sum of exponentials of input elements.
 
@@ -58,7 +60,8 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
 
     See Also
     --------
-    numpy.logaddexp, numpy.logaddexp2
+    :data:`numpy.logaddexp`
+    :data:`numpy.logaddexp2`
 
     Notes
     -----
@@ -246,6 +249,7 @@ def _logsumexp(a, b, *, axis, return_sign, xp):
     return out, sgn
 
 
+@xp_capabilities()
 def softmax(x, axis=None):
     r"""Compute the softmax function.
 
@@ -344,6 +348,7 @@ def softmax(x, axis=None):
     return exp_x_shifted / xp.sum(exp_x_shifted, axis=axis, keepdims=True)
 
 
+@xp_capabilities()
 def log_softmax(x, axis=None):
     r"""Compute the logarithm of the softmax function.
 
diff --git a/scipy/special/tests/test_logsumexp.py b/scipy/special/tests/test_logsumexp.py
@@ -4,7 +4,8 @@
 
 import numpy as np
 
-from scipy._lib._array_api import is_array_api_strict, xp_default_dtype, xp_device
+from scipy._lib._array_api import (is_array_api_strict, make_skip_xp_backends,
+                                   xp_default_dtype, xp_device)
 from scipy._lib._array_api_no_0d import (xp_assert_equal, xp_assert_close,
                                          xp_assert_less)
 
@@ -38,6 +39,7 @@ def test_wrap_radians(xp):
 @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning")
 @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning")
 @pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning")
+@make_skip_xp_backends(logsumexp)
 class TestLogSumExp:
     def test_logsumexp(self, xp):
         # Test with zero-size array
@@ -310,6 +312,7 @@ def test_device(self, x_raw, xp, devices):
             assert xp_device(logsumexp(x, b=x)) == xp_device(x)
 
 
+@make_skip_xp_backends(softmax)
 class TestSoftmax:
     def test_softmax_fixtures(self, xp):
         xp_assert_close(softmax(xp.asarray([1000., 0., 0., 0.])),
@@ -378,6 +381,7 @@ def test_softmax_array_like(self):
                         np.asarray([1., 0., 0., 0.]), rtol=1e-13)
 
 
+@make_skip_xp_backends(log_softmax)
 class TestLogSoftmax:
     def test_log_softmax_basic(self, xp):
         xp_assert_close(log_softmax(xp.asarray([1000., 1.])),
