Repository: scipy/scipy
Pull request #22946: ENH: `special`: add `xp_capabilities`
Author: login: crusaderky
Labels: name: enhancement
description: A new feature or improvement
name: scipy.special
name: scipy._lib
name: array types
description: Items related to array API support and input array validation (see gh-18286)
Created: 2025-05-08T11:51:31Z | Merged: 2025-05-15T18:05:15Z

- Add `@xp_capabilities` to the ufuncs in `scipy.special._support_alternative_backends`
- Speed up function dispatch by ~4us by applying the same caching pattern already used in array-api-compat helpers
- Fix bug where the wrapped functions could not be pickled
- Fix bug where the wrapped functions would have poor `repr()`, `__qualname__`, and `__module__`

### Demo render
![image](https://github.com/user-attachments/assets/9129d0e1-4bb3-4872-b9fe-32b5baa6fca7)



Diff:
diff --git a/scipy/_lib/_array_api.py b/scipy/_lib/_array_api.py
index 8641b7ded8c9..aa9d9c524af2 100644
--- a/scipy/_lib/_array_api.py
+++ b/scipy/_lib/_array_api.py
@@ -798,17 +798,21 @@ def xp_capabilities(
     sphinx_capabilities = _make_sphinx_capabilities(**capabilities)
 
     def decorator(f):
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            return f(*args, **kwargs)
-
-        capabilities_table[wrapper] = capabilities
+        # Don't use a wrapper, as in some cases @xp_capabilities is
+        # applied to a ufunc
+        capabilities_table[f] = capabilities
         note = _make_capabilities_note(f.__name__, sphinx_capabilities)
-        doc = FunctionDoc(wrapper)
+        doc = FunctionDoc(f)
         doc['Notes'].append(note)
-        wrapper.__doc__ = str(doc).split("\n", 1)[1]  # remove signature
+        doc = str(doc).split("\n", 1)[1]  # remove signature
+        try:
+            f.__doc__ = doc
+        except AttributeError:
+            # Can't update __doc__ on ufuncs if SciPy
+            # was compiled against NumPy < 2.2.
+            pass
 
-        return wrapper
+        return f
     return decorator
 
 
@@ -868,7 +872,7 @@ def make_xp_test_case(*funcs, capabilities_table=None):
     return lambda func: functools.reduce(lambda f, g: g(f), marks, func)
 
 
-def make_xp_pytest_param(func, capabilities_table=None):
+def make_xp_pytest_param(func, *args, capabilities_table=None):
     """Variant of ``make_xp_test_case`` that returns a pytest.param for a function,
     with all necessary skip_xp_backends and xfail_xp_backends marks applied::
     
@@ -892,6 +896,20 @@ def test(func, xp):
         def test(func, xp):
             ...
 
+    Parameters
+    ----------
+    func : Callable
+        Function to be tested. It must be decorated with ``@xp_capabilities``.
+    *args : Any, optional
+        Extra pytest parameters for the use case, e.g.::
+
+        @pytest.mark.parametrize("func,verb", [
+            make_xp_pytest_param(f1, "hello"),
+            make_xp_pytest_param(f2, "world")])
+        def test(func, verb, xp):
+            # iterates on (func=f1, verb="hello")
+            # and (func=f2, verb="world")
+
     See Also
     --------
     xp_capabilities
@@ -901,7 +919,7 @@ def test(func, xp):
     import pytest
 
     marks = _make_xp_pytest_marks(func, capabilities_table=capabilities_table)
-    return pytest.param(func, marks=marks)
+    return pytest.param(func, *args, marks=marks, id=func.__name__)
 
 
 # Is it OK to have a dictionary that is mutated (once upon import) in many places?
diff --git a/scipy/special/__init__.py b/scipy/special/__init__.py
index 0c2fe8809279..2572e27b48d8 100644
--- a/scipy/special/__init__.py
+++ b/scipy/special/__init__.py
@@ -785,10 +785,7 @@
 from ._ufuncs import *
 
 # Replace some function definitions from _ufuncs to add Array API support
-from ._support_alternative_backends import (
-    log_ndtr, ndtr, ndtri, erf, erfc, i0, i0e, i1, i1e, gammaln,
-    gammainc, gammaincc, logit, expit, entr, rel_entr, xlogy,
-    chdtr, chdtrc, betainc, betaincc, stdtr, stdtrit)
+from ._support_alternative_backends import *
 
 from . import _basic
 from ._basic import *
diff --git a/scipy/special/_support_alternative_backends.py b/scipy/special/_support_alternative_backends.py
index 74a1c4f32418..e9811fcfe495 100644
--- a/scipy/special/_support_alternative_backends.py
+++ b/scipy/special/_support_alternative_backends.py
@@ -1,65 +1,142 @@
 import functools
 import operator
+from collections.abc import Callable
+from dataclasses import dataclass
+from types import ModuleType
 
 import numpy as np
 from scipy._lib._array_api import (
     array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray,
-    xp_promote, SCIPY_ARRAY_API
+    xp_promote, xp_capabilities, SCIPY_ARRAY_API
 )
 import scipy._lib.array_api_extra as xpx
 from . import _ufuncs
-# These don't really need to be imported, but otherwise IDEs might not realize
-# that these are defined in this file / report an error in __init__.py
-from ._ufuncs import (
-    log_ndtr, ndtr, ndtri, erf, erfc, i0, i0e, i1, i1e, gammaln,  # noqa: F401
-    gammainc, gammaincc, logit, expit, entr, rel_entr, xlogy,  # noqa: F401
-    chdtr, chdtrc, betainc, betaincc, stdtr, stdtrit  # noqa: F401
-)
-
-array_api_compat_prefix = "scipy._lib.array_api_compat"
 
 
-def get_array_special_func(f_name, xp):
-    if is_numpy(xp):
-        return getattr(_ufuncs, f_name)
-
-    spx = scipy_namespace_for(xp)
-    if spx is not None:
-        f = getattr(spx.special, f_name, None)
+@dataclass
+class _FuncInfo:
+    # NumPy-only function. IT MUST BE ELEMENTWISE.
+    func: Callable
+    # Number of arguments, not counting out=
+    # This is for testing purposes only, due to the fact that
+    # inspect.signature() just returns *args for ufuncs.
+    n_args: int
+    # @xp_capabilities decorator, for the purpose of
+    # documentation and unit testing. Omit to indicate
+    # full support for all backends.
+    xp_capabilities: Callable[[Callable], Callable] | None = None
+    # Generic implementation to fall back on if there is no native dispatch
+    # available. This is a function that accepts (main namespace, scipy namespace)
+    # and returns the final callable, or None if not available.
+    generic_impl: Callable[
+        [ModuleType, ModuleType | None], Callable | None
+    ] | None = None
+
+    @property
+    def name(self):
+        return self.func.__name__
+
+    # These are needed by @lru_cache below
+    def __hash__(self):
+        return hash(self.func)
+
+    def __eq__(self, other):
+        return isinstance(other, _FuncInfo) and self.func == other.func
+
+    @property
+    def wrapper(self):
+        if self.name in globals():
+            # Already initialised. We are likely in a unit test.
+            # Return function potentially overridden by xpx.testing.lazy_xp_function.
+            import scipy.special
+            return getattr(scipy.special, self.name)
+
+        if SCIPY_ARRAY_API:
+            @functools.wraps(self.func)
+            def wrapped(*args, **kwargs):
+                xp = array_namespace(*args)
+                return self._wrapper_for(xp)(*args, **kwargs)
+
+            # Allow pickling the function. Normally this is done by @wraps,
+            # but in this case it doesn't work because self.func is a ufunc.
+            wrapped.__module__ = "scipy.special"
+            wrapped.__qualname__ = self.name
+            func = wrapped
+        else:
+            func = self.func
+
+        capabilities = self.xp_capabilities or xp_capabilities()
+        # In order to retain a naked ufunc when SCIPY_ARRAY_API is
+        # disabled, xp_capabilities must apply its changes in place.
+        cap_func = capabilities(func)
+        assert cap_func is func
+        return func
+
+    @functools.lru_cache(1000)
+    def _wrapper_for(self, xp):
+        if is_numpy(xp):
+            return self.func
+
+        # If a native implementation is available, use that
+        spx = scipy_namespace_for(xp)
+        f = _get_native_func(xp, spx, self.name)
         if f is not None:
             return f
 
-    # if generic array-API implementation is available, use that;
-    # otherwise, fall back to NumPy/SciPy
-    if f_name in _generic_implementations:
-        f = _generic_implementations[f_name](xp=xp, spx=spx)
-        if f is not None:
-            return f
+        # If generic Array API implementation is available, use that
+        if self.generic_impl is not None:
+            f = self.generic_impl(xp, spx)
+            if f is not None:
+                return f
 
-    def f(*args, **kwargs):
         if is_marray(xp):
-            _f = globals()[f_name]  # Allow nested wrapping
-            data_args = [arg.data for arg in args]
-            out = _f(*data_args, **kwargs)
-            mask = functools.reduce(operator.or_, (arg.mask for arg in args))
-            return xp.asarray(out, mask=mask)
-
-        elif is_dask(xp):
-            # IMPORTANT: map_blocks works only because all ufuncs in this module
+            # Unwrap the array, apply the function on the wrapped namespace,
+            # and then re-wrap it.
+            # IMPORTANT: this only works because all functions in this module
+            # are elementwise. Otherwise, we would not be able to define a
+            # general rule for mask propagation.
+
+            _f = globals()[self.name]  # Allow nested wrapping
+            def f(*args, _f=_f, xp=xp, **kwargs):
+                data_args = [arg.data for arg in args]
+                out = _f(*data_args, **kwargs)
+                mask = functools.reduce(operator.or_, (arg.mask for arg in args))
+                return xp.asarray(out, mask=mask)
+
+            return f
+
+        if is_dask(xp):
+            # Apply the function to each block of the Dask array.
+            # IMPORTANT: map_blocks works only because all functions in this module
             # are elementwise. It would be a grave mistake to apply this to gufuncs
             # or any other function with reductions, as they would change their
             # output depending on chunking!
 
-            _f = globals()[f_name]  # Allow nested wrapping
-            # Hide dtype kwarg from map_blocks
-            return xp.map_blocks(functools.partial(_f, **kwargs), *args)
+            _f = globals()[self.name]  # Allow nested wrapping
+            def f(*args, _f=_f, xp=xp, **kwargs):
+                # Hide dtype kwarg from map_blocks
+                return xp.map_blocks(functools.partial(_f, **kwargs), *args)
 
-        else:
-            _f = getattr(_ufuncs, f_name)
+            return f
+
+        # As a final resort, use the NumPy/SciPy implementation
+        _f = self.func
+        def f(*args, _f=_f, xp=xp, **kwargs):
+            # TODO use xpx.lazy_apply to add jax.jit support
+            # (but dtype propagation can be non-trivial)
             args = [np.asarray(arg) for arg in args]
             out = _f(*args, **kwargs)
             return xp.asarray(out)
 
+        return f
+
+
+def _get_native_func(xp, spx, f_name):
+    f = getattr(spx.special, f_name, None) if spx else None
+    if f is None and hasattr(xp, 'special'):
+        # Currently dead branch, in anticipation of 'special' Array API extension
+        # https://github.com/data-apis/array-api/issues/725
+        f = getattr(xp.special, f_name, None)
     return f
 
 
@@ -93,19 +170,13 @@ def __xlogy(x, y, *, xp=xp):
     return __xlogy
 
 
-def _get_native_func(xp, spx, f_name):
-    f = getattr(spx.special, f_name, None) if spx else None
-    if f is None and hasattr(xp, 'special'):
-        f = getattr(xp.special, f_name, None)
-    return f
-
 
 def _chdtr(xp, spx):
     # The difference between this and just using `gammainc`
     # defined by `get_array_special_func` is that if `gammainc`
     # isn't found, we don't want to use the SciPy version; we'll
     # return None here and use the SciPy version of `chdtr`.
-    gammainc = _get_native_func(xp, spx, 'gammainc')  # noqa: F811
+    gammainc = _get_native_func(xp, spx, 'gammainc')
     if gammainc is None:
         return None
 
@@ -124,7 +195,7 @@ def _chdtrc(xp, spx):
     # defined by `get_array_special_func` is that if `gammaincc`
     # isn't found, we don't want to use the SciPy version; we'll
     # return None here and use the SciPy version of `chdtrc`.
-    gammaincc = _get_native_func(xp, spx, 'gammaincc')  # noqa: F811
+    gammaincc = _get_native_func(xp, spx, 'gammaincc')
     if gammaincc is None:
         return None
 
@@ -137,7 +208,7 @@ def __chdtrc(v, x):
 
 
 def _betaincc(xp, spx):
-    betainc = _get_native_func(xp, spx, 'betainc')  # noqa: F811
+    betainc = _get_native_func(xp, spx, 'betainc')
     if betainc is None:
         return None
 
@@ -148,7 +219,7 @@ def __betaincc(a, b, x):
 
 
 def _stdtr(xp, spx):
-    betainc = _get_native_func(xp, spx, 'betainc')  # noqa: F811
+    betainc = _get_native_func(xp, spx, 'betainc')
     if betainc is None:
         return None
 
@@ -161,11 +232,12 @@ def __stdtr(df, t):
 
 
 def _stdtrit(xp, spx):
-    betainc = _get_native_func(xp, spx, 'betainc')  # noqa: F811
+    # Need either native stdtr or native betainc
+    stdtr = _get_native_func(xp, spx, 'stdtr') or _stdtr(xp, spx)
     # If betainc is not defined, the root-finding would be done with `xp`
     # despite `stdtr` being evaluated with SciPy/NumPy `stdtr`. Save the
     # conversions: in this case, just evaluate `stdtrit` with SciPy/NumPy.
-    if betainc is None:
+    if stdtr is None:
         return None
 
     from scipy.optimize.elementwise import bracket_root, find_root
@@ -179,62 +251,45 @@ def fun(t, df, p):  return stdtr(df, t) - p
     return __stdtrit
 
 
-_generic_implementations = {'rel_entr': _rel_entr,
-                            'xlogy': _xlogy,
-                            'chdtr': _chdtr,
-                            'chdtrc': _chdtrc,
-                            'betaincc': _betaincc,
-                            'stdtr': _stdtr,
-                            'stdtrit': _stdtrit,
-                            }
-
-
-# functools.wraps doesn't work because:
-# 'numpy.ufunc' object has no attribute '__module__'
-def support_alternative_backends(f_name):
-    func = getattr(_ufuncs, f_name)
-
-    @functools.wraps(func)
-    def wrapped(*args, **kwargs):
-        xp = array_namespace(*args)
-        f = get_array_special_func(f_name, xp)
-        return f(*args, **kwargs)
-
-    return wrapped
-
-
-# function name: number of args (for testing purposes)
-array_special_func_map = {
-    'log_ndtr': 1,
-    'ndtr': 1,
-    'ndtri': 1,
-    'erf': 1,
-    'erfc': 1,
-    'i0': 1,
-    'i0e': 1,
-    'i1': 1,
-    'i1e': 1,
-    'gammaln': 1,
-    'gammainc': 2,
-    'gammaincc': 2,
-    'logit': 1,
-    'expit': 1,
-    'entr': 1,
-    'rel_entr': 2,
-    'xlogy': 2,
-    'chdtr': 2,
-    'chdtrc': 2,
-    'betainc': 3,
-    'betaincc': 3,
-    'stdtr': 2,
-    'stdtrit': 2,
-}
-
-globals().update(
-    {f_name: support_alternative_backends(f_name)
-     if SCIPY_ARRAY_API
-     else getattr(_ufuncs, f_name)
-     for f_name in array_special_func_map}
+# Inventory of automatically dispatched functions
+# IMPORTANT: these must all be **elementwise** functions!
+
+# PyTorch doesn't implement `betainc`.
+# On torch CPU we can fall back to NumPy, but on GPU it won't work.
+_needs_betainc = xp_capabilities(cpu_only=True, exceptions=['jax.numpy', 'cupy'])
+
+_special_funcs = (
+    _FuncInfo(_ufuncs.betainc, 3, _needs_betainc),
+    _FuncInfo(_ufuncs.betaincc, 3, _needs_betainc, generic_impl=_betaincc),
+    _FuncInfo(_ufuncs.chdtr, 2, generic_impl=_chdtr),
+    _FuncInfo(_ufuncs.chdtrc, 2, generic_impl=_chdtrc),
+    _FuncInfo(_ufuncs.erf, 1),
+    _FuncInfo(_ufuncs.erfc, 1),
+    _FuncInfo(_ufuncs.entr, 1),
+    _FuncInfo(_ufuncs.expit, 1),
+    _FuncInfo(_ufuncs.i0, 1),
+    _FuncInfo(_ufuncs.i0e, 1),
+    _FuncInfo(_ufuncs.i1, 1),
+    _FuncInfo(_ufuncs.i1e, 1),
+    _FuncInfo(_ufuncs.log_ndtr, 1),
+    _FuncInfo(_ufuncs.logit, 1),
+    _FuncInfo(_ufuncs.gammaln, 1),
+    _FuncInfo(_ufuncs.gammainc, 2),
+    _FuncInfo(_ufuncs.gammaincc, 2),
+    _FuncInfo(_ufuncs.ndtr, 1),
+    _FuncInfo(_ufuncs.ndtri, 1),
+    _FuncInfo(_ufuncs.rel_entr, 2, generic_impl=_rel_entr),
+    _FuncInfo(_ufuncs.stdtr,  2, _needs_betainc, generic_impl=_stdtr),
+    _FuncInfo(_ufuncs.stdtrit, 2,
+              xp_capabilities(
+                  cpu_only=True, exceptions=['cupy'],  # needs betainc
+                  skip_backends=[("jax.numpy", "no scipy.optimize support")]),
+              generic_impl=_stdtrit),
+    _FuncInfo(_ufuncs.xlogy, 2, generic_impl=_xlogy),
 )
 
-__all__ = list(array_special_func_map)
+# Override ufuncs.
+# When SCIPY_ARRAY_API is disabled, this exclusively updates the docstrings in place
+# and populates the xp_capabilities table, while retaining the original ufuncs.
+globals().update({nfo.func.__name__: nfo.wrapper for nfo in _special_funcs})
+__all__ = [nfo.func.__name__ for nfo in _special_funcs]
diff --git a/scipy/special/tests/test_support_alternative_backends.py b/scipy/special/tests/test_support_alternative_backends.py
index 78e20046df9d..e98939fef4b1 100644
--- a/scipy/special/tests/test_support_alternative_backends.py
+++ b/scipy/special/tests/test_support_alternative_backends.py
@@ -1,35 +1,24 @@
 from functools import partial
-from types import ModuleType
+import pickle
+
 import pytest
 from hypothesis import given, strategies
 import hypothesis.extra.numpy as npst
+from packaging import version
 
 from scipy import special
-from scipy.special._support_alternative_backends import (get_array_special_func,
-                                                         array_special_func_map)
+from scipy.special._support_alternative_backends import _special_funcs
 from scipy._lib._array_api_no_0d import xp_assert_close
 from scipy._lib._array_api import (is_cupy, is_dask, is_jax, is_torch,
-                                   xp_default_dtype, SCIPY_ARRAY_API, SCIPY_DEVICE)
+                                   make_xp_pytest_param, make_xp_test_case,
+                                   xp_default_dtype)
 from scipy._lib.array_api_compat import numpy as np
-from scipy._lib.array_api_extra.testing import lazy_xp_function
-
-
-special_wrapped = ModuleType("special_wrapped")
-lazy_xp_modules = [special_wrapped]
-for f_name in array_special_func_map:
-    f = getattr(special, f_name)
-    setattr(special_wrapped, f_name, f)
-    lazy_xp_function(f)
 
+# Run all tests in this module in the Array API CI, including those without
+# the xp fixture
+pytestmark = pytest.mark.array_api_backends
 
-@pytest.mark.skipif(not SCIPY_ARRAY_API, reason="Alternative backends must be enabled.")
-def test_dispatch_to_unrecognized_library():
-    xp = pytest.importorskip("array_api_strict")
-    f = get_array_special_func('ndtr', xp=xp)
-    x = [1, 2, 3]
-    res = f(xp.asarray(x))
-    ref = xp.asarray(special.ndtr(np.asarray(x)))
-    xp_assert_close(res, ref)
+lazy_xp_modules = [special]
 
 
 def _skip_or_tweak_alternative_backends(xp, f_name, dtypes):
@@ -44,15 +33,6 @@ def _skip_or_tweak_alternative_backends(xp, f_name, dtypes):
     dtypes_np_ref : list[str]
         The dtypes to use for the reference NumPy arrays.
     """
-    if (SCIPY_DEVICE != 'cpu'
-        and is_torch(xp)
-        and f_name in {'stdtr', 'stdtrit', 'betaincc', 'betainc'}
-    ):
-        pytest.skip(f"`{f_name}` does not have an array-agnostic implementation "
-                    "and cannot delegate to PyTorch.")
-    if is_jax(xp) and f_name == "stdtrit":
-        pytest.skip(f"`{f_name}` requires scipy.optimize support for immutable arrays")
-
     if ((is_jax(xp) and f_name == 'gammaincc')  # google/jax#20699
         # gh-20972
         or ((is_cupy(xp) or is_jax(xp) or is_torch(xp)) and f_name == 'chdtrc')):
@@ -93,20 +73,19 @@ def _skip_or_tweak_alternative_backends(xp, f_name, dtypes):
     return positive_only, dtypes_np_ref
 
 
-@pytest.mark.parametrize('f_name,n_args', array_special_func_map.items())
 @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask")
-@pytest.mark.parametrize('dtype', ['float32', 'float64', 'int64'])
 @pytest.mark.parametrize('shapes', [[(0,)]*4, [tuple()]*4, [(10,)]*4,
                                     [(10,), (11, 1), (12, 1, 1), (13, 1, 1, 1)]])
-def test_support_alternative_backends(xp, f_name, n_args, dtype, shapes):
+@pytest.mark.parametrize('dtype', ['float32', 'float64', 'int64'])
+@pytest.mark.parametrize(
+    'func,nfo', [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs])
+def test_support_alternative_backends(xp, func, nfo, dtype, shapes):
     positive_only, [dtype_np_ref] = _skip_or_tweak_alternative_backends(
-        xp, f_name, [dtype])
-    f = getattr(special, f_name)  # Unwrapped
-    fw = getattr(special_wrapped, f_name)  # Wrapped by lazy_xp_function
+        xp, nfo.name, [dtype])
     dtype_np = getattr(np, dtype)
     dtype_xp = getattr(xp, dtype)
 
-    shapes = shapes[:n_args]
+    shapes = shapes[:nfo.n_args]
     rng = np.random.default_rng(984254252920492019)
     if 'int' in dtype:
         iinfo = np.iinfo(dtype_np)
@@ -127,27 +106,24 @@ def test_support_alternative_backends(xp, f_name, n_args, dtype, shapes):
         # Try to trigger bugs related to having multiple chunks.
         args_xp = [arg.rechunk(5) for arg in args_xp]
 
-    res = fw(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped ufunc
 
     # When dtype_np is integer, the output dtype can be float
     atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps
     xp_assert_close(res, xp.asarray(ref), atol=atol)
 
 
-@pytest.mark.parametrize('f_name,n_args',
-                         [(f_name, n_args)
-                          for f_name, n_args in array_special_func_map.items()
-                          if n_args >= 2])
+@pytest.mark.parametrize(
+    'func, nfo',
+    [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs if i.n_args >= 2])
 @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask")
-def test_support_alternative_backends_mismatched_dtypes(xp, f_name, n_args):
+def test_support_alternative_backends_mismatched_dtypes(xp, func, nfo):
     """Test mix-n-match of int and float arguments"""
-    assert n_args <= 3
-    dtypes = ['int64', 'float32', 'float64'][:n_args]
-    dtypes_xp = [xp.int64, xp.float32, xp.float64][:n_args]
+    dtypes = ['int64', 'float32', 'float64'][:nfo.n_args]
+    dtypes_xp = [xp.int64, xp.float32, xp.float64][:nfo.n_args]
     positive_only, dtypes_np_ref = _skip_or_tweak_alternative_backends(
-        xp, f_name, dtypes)
-    f = getattr(special, f_name)
+        xp, nfo.name, dtypes)
 
     rng = np.random.default_rng(984254252920492019)
     iinfo = np.iinfo(np.int64)
@@ -156,7 +132,7 @@ def test_support_alternative_backends_mismatched_dtypes(xp, f_name, n_args):
         randint(size=1, dtype=np.int64),
         rng.standard_normal(size=1, dtype=np.float32),
         rng.standard_normal(size=1, dtype=np.float64),
-    ][:n_args]
+    ][:nfo.n_args]
     if positive_only:
         args_np = [np.abs(arg) for arg in args_np]
 
@@ -165,8 +141,8 @@ def test_support_alternative_backends_mismatched_dtypes(xp, f_name, n_args):
     args_np = [np.asarray(arg, dtype=dtype_np_ref)
                for arg, dtype_np_ref in zip(args_np, dtypes_np_ref)]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped ufunc
 
     atol = 10 * np.finfo(ref.dtype).eps
     xp_assert_close(res, xp.asarray(ref), atol=atol)
@@ -175,18 +151,18 @@ def test_support_alternative_backends_mismatched_dtypes(xp, f_name, n_args):
 @pytest.mark.xslow
 @given(data=strategies.data())
 @pytest.mark.fail_slow(5)
-# `reversed` is for developer convenience: test new function first = less waiting
-@pytest.mark.parametrize('f_name,n_args', reversed(array_special_func_map.items()))
+@pytest.mark.parametrize(
+    'func,nfo', [make_xp_pytest_param(nfo.wrapper, nfo) for nfo in _special_funcs])
 @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask")
 @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask")
+@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning:dask")
 @pytest.mark.filterwarnings(
     "ignore:overflow encountered:RuntimeWarning:array_api_strict"
 )
-def test_support_alternative_backends_hypothesis(xp, f_name, n_args, data):
+def test_support_alternative_backends_hypothesis(xp, func, nfo, data):
     dtype = data.draw(strategies.sampled_from(['float32', 'float64', 'int64']))
     positive_only, [dtype_np_ref] = _skip_or_tweak_alternative_backends(
-        xp, f_name, [dtype])
-    f = getattr(special, f_name)
+        xp, nfo.name, [dtype])
     dtype_np = getattr(np, dtype)
     dtype_xp = getattr(xp, dtype)
 
@@ -197,21 +173,70 @@ def test_support_alternative_backends_hypothesis(xp, f_name, n_args, data):
     if positive_only:
         elements['min_value'] = 0
 
-    shapes, _ = data.draw(npst.mutually_broadcastable_shapes(num_shapes=n_args))
+    shapes, _ = data.draw(
+        npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args))
     args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements))
                for shape in shapes]
 
     args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np]
     args_np = [np.asarray(arg, dtype=dtype_np_ref) for arg in args_np]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped ufunc
 
     # When dtype_np is integer, the output dtype can be float
     atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps
     xp_assert_close(res, xp.asarray(ref), atol=atol)
 
 
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_pickle(func):
+    roundtrip = pickle.loads(pickle.dumps(func))
+    assert roundtrip is func
+
+
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_repr(func):
+    assert func.__name__ in repr(func)
+    assert "locals" not in repr(func)
+
+
+@pytest.mark.skipif(
+    version.parse(np.__version__) < version.parse("2.2"),
+    reason="Can't update ufunc __doc__ when SciPy is compiled vs. NumPy < 2.2")
+@pytest.mark.parametrize('func', [nfo.wrapper for nfo in _special_funcs])
+def test_doc(func):
+    """xp_capabilities updates the docstring in place. 
+    Make sure it does so exactly once, including when SCIPY_ARRAY_API is not set.
+    """
+    match = "has experimental support for Python Array API"
+    assert func.__doc__.count(match) == 1
+
+
+@pytest.mark.parametrize('func,n_args',
+                         [(nfo.wrapper, nfo.n_args) for nfo in _special_funcs])
+def test_ufunc_kwargs(func, n_args):
+    """Test that numpy-specific out= and dtype= keyword arguments
+    of ufuncs still work when SCIPY_ARRAY_API is set.
+    """
+    # out=
+    args = [np.asarray([.1, .2])] * n_args
+    out = np.empty(2)
+    y = func(*args, out=out)
+    xp_assert_close(y, out)
+
+    # out= with out.dtype != args.dtype
+    out = np.empty(2, dtype=np.float32)
+    y = func(*args, out=out)
+    xp_assert_close(y, out)
+
+    # dtype=
+    y = func(*args, dtype=np.float32)
+    assert y.dtype == np.float32
+
+
+
+@make_xp_test_case(special.chdtr)
 def test_chdtr_gh21311(xp):
     # the edge case behavior of generic chdtr was not right; see gh-21311
     # be sure to test at least these cases


Reviews:
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T11:55:38Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T12:06:33Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T12:06:48Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T12:09:14Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T12:23:10Z
user: login: steppi
body: Very nice work Guido. Just a few minor comments.
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-15T14:57:21Z
user: login: steppi
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T15:06:13Z
user: login: steppi
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T15:06:40Z
user: login: steppi
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T15:07:00Z
user: login: steppi
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-15T15:07:18Z
user: login: crusaderky
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T15:11:10Z
user: login: steppi
body: Great! I'll let CI settle, but this looks like it's ready to go in.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-15T15:14:05Z

Review comments:
diff_hunk: @@ -93,19 +165,13 @@ def __xlogy(x, y, *, xp=xp):
     return __xlogy
 
 
-def _get_native_func(xp, spx, f_name):
-    f = getattr(spx.special, f_name, None) if spx else None
-    if f is None and hasattr(xp, 'special'):
-        f = getattr(xp.special, f_name, None)
path: scipy/special/_support_alternative_backends.py
user: login: crusaderky
body: Removed these two lines. There is no use case for them.
created_at: 2025-05-08T11:55:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,35 +1,23 @@
 from functools import partial
-from types import ModuleType
+import pickle
+
 import pytest
 from hypothesis import given, strategies
 import hypothesis.extra.numpy as npst
 
 from scipy import special
-from scipy.special._support_alternative_backends import (get_array_special_func,
-                                                         array_special_func_map)
+from scipy.special._support_alternative_backends import _special_funcs
 from scipy._lib._array_api_no_0d import xp_assert_close
 from scipy._lib._array_api import (is_cupy, is_dask, is_jax, is_torch,
-                                   xp_default_dtype, SCIPY_ARRAY_API, SCIPY_DEVICE)
+                                   make_xp_pytest_param, make_skip_xp_backends,
+                                   xp_default_dtype)
 from scipy._lib.array_api_compat import numpy as np
-from scipy._lib.array_api_extra.testing import lazy_xp_function
-
-
-special_wrapped = ModuleType("special_wrapped")
-lazy_xp_modules = [special_wrapped]
-for f_name in array_special_func_map:
-    f = getattr(special, f_name)
-    setattr(special_wrapped, f_name, f)
-    lazy_xp_function(f)
 
+# Run all tests in this module in the Array API CI, including those without
+# the xp fixture
+pytestmark = pytest.mark.array_api_backends
 
-@pytest.mark.skipif(not SCIPY_ARRAY_API, reason="Alternative backends must be enabled.")
-def test_dispatch_to_unrecognized_library():
path: scipy/special/tests/test_support_alternative_backends.py
user: login: crusaderky
body: Unneeded - _all_ tests below that run on array-api-strict already test this.
created_at: 2025-05-08T12:06:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,15 +32,6 @@ def _skip_or_tweak_alternative_backends(xp, f_name, dtypes):
     dtypes_np_ref : list[str]
         The dtypes to use for the reference NumPy arrays.
     """
-    if (SCIPY_DEVICE != 'cpu'
-        and is_torch(xp)
-        and f_name in {'stdtr', 'stdtrit', 'betaincc', 'betainc'}
-    ):
-        pytest.skip(f"`{f_name}` does not have an array-agnostic implementation "
-                    "and cannot delegate to PyTorch.")
-    if is_jax(xp) and f_name == "stdtrit":
-        pytest.skip(f"`{f_name}` requires scipy.optimize support for immutable arrays")
path: scipy/special/tests/test_support_alternative_backends.py
user: login: crusaderky
body: Above: moved to xp_capabilities.
Below: these special cases are too niche for xp_capabilities so they stay here.
created_at: 2025-05-08T12:06:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -127,27 +104,23 @@ def test_support_alternative_backends(xp, f_name, n_args, dtype, shapes):
         # Try to trigger bugs related to having multiple chunks.
         args_xp = [arg.rechunk(5) for arg in args_xp]
 
-    res = fw(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
path: scipy/special/tests/test_support_alternative_backends.py
user: login: crusaderky
body: To test `lazy_xp_function`:
1. go to `xlogy`, add the line `assert xp.any(x) or ~xp.any(x)`, and watch Dask fail with spurious compute error
2. go to `stdtr`, add the line `assert xp.any(t) or ~xp.any(t)`, and watch JAX fail with ConcretizationError
created_at: 2025-05-08T12:09:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -161,11 +227,12 @@ def __stdtr(df, t):
 
 
 def _stdtrit(xp, spx):
-    betainc = _get_native_func(xp, spx, 'betainc')  # noqa: F811
+    # Need either native stdtr or native betainc
+    stdtr = _get_native_func(spx, 'stdtr') or _stdtr(xp, spx)
     # If betainc is not defined, the root-finding would be done with `xp`
     # despite `stdtr` being evaluated with SciPy/NumPy `stdtr`. Save the
     # conversions: in this case, just evaluate `stdtrit` with SciPy/NumPy.
-    if betainc is None:
+    if stdtr is None:
path: scipy/special/_support_alternative_backends.py
user: login: crusaderky
body: Note: this is purely for the sake of readability. There is no actual use case where a backend defines stdtr but not betainc.
created_at: 2025-05-08T12:23:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -789,17 +789,21 @@ def xp_capabilities(
     sphinx_capabilities = _make_sphinx_capabilities(**capabilities)
 
     def decorator(f):
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            return f(*args, **kwargs)
-
-        capabilities_table[wrapper] = capabilities
+        # Don't use a wrapper, as in some cases @xp_capabilities is
+        # applied to a ufunc
+        capabilities_table[f] = capabilities
path: scipy/_lib/_array_api.py
user: login: steppi
body:  I took a look at https://github.com/scipy/scipy/pull/22686 first, and thought to make this suggestion. Great to see it's already done here.
created_at: 2025-05-14T20:02:49Z
author_association: MEMBER
diff_hunk: @@ -789,17 +789,21 @@ def xp_capabilities(
     sphinx_capabilities = _make_sphinx_capabilities(**capabilities)
 
     def decorator(f):
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            return f(*args, **kwargs)
-
-        capabilities_table[wrapper] = capabilities
+        # Don't use a wrapper, as in some cases @xp_capabilities is
+        # applied to a ufunc
+        capabilities_table[f] = capabilities
         note = _make_capabilities_note(f.__name__, sphinx_capabilities)
-        doc = FunctionDoc(wrapper)
+        doc = FunctionDoc(f)
         doc['Notes'].append(note)
-        wrapper.__doc__ = str(doc).split("\n", 1)[1]  # remove signature
+        doc = str(doc).split("\n", 1)[1]  # remove signature
+        try:
+            f.__doc__ = doc
+        except AttributeError:
+            # Can't update __doc__ on Cython ufuncs if SciPy
+            # was compiled against NumPy 1.x
path: scipy/_lib/_array_api.py
user: login: steppi
body: ```suggestion
            # Can't update __doc__ on ufuncs if SciPy
            # was compiled against NumPy < 2.2.
```

It actually wasn't possible to update `__doc__` for any ufuncs prior to NumPy 2.2. See https://github.com/numpy/numpy/pull/27735. Also, ufuncs are a fundamental primitive in NumPy and this doesn't have anything to do with Cython.
created_at: 2025-05-14T21:05:43Z
author_association: MEMBER
diff_hunk: @@ -1,66 +1,138 @@
 import functools
 import operator
+from collections.abc import Callable
+from dataclasses import dataclass
+from types import ModuleType
 
 import numpy as np
 from scipy._lib._array_api import (
     array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray,
-    xp_promote, SCIPY_ARRAY_API
+    xp_promote, xp_capabilities, SCIPY_ARRAY_API
 )
 import scipy._lib.array_api_extra as xpx
 from . import _ufuncs
-# These don't really need to be imported, but otherwise IDEs might not realize
-# that these are defined in this file / report an error in __init__.py
-from ._ufuncs import (
-    log_ndtr, ndtr, ndtri, erf, erfc, i0, i0e, i1, i1e, gammaln,  # noqa: F401
-    gammainc, gammaincc, logit, expit, entr, rel_entr, xlogy,  # noqa: F401
-    chdtr, chdtrc, betainc, betaincc, stdtr, stdtrit  # noqa: F401
-)
-
-array_api_compat_prefix = "scipy._lib.array_api_compat"
 
 
-def get_array_special_func(f_name, xp):
-    if is_numpy(xp):
-        return getattr(_ufuncs, f_name)
-
-    spx = scipy_namespace_for(xp)
-    if spx is not None:
-        f = getattr(spx.special, f_name, None)
+@dataclass
+class _FuncInfo:
+    # NumPy-only function. IT MUST BE ELEMENTWISE.
+    func: Callable
+    # Number of arguments, not counting out=
+    # This is for testing purposes only, due to the fact that
+    # inspect.signature() just returns *args for Cython ufuncs.
path: scipy/special/_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    # inspect.signature() just returns *args for ufuncs.
```
created_at: 2025-05-14T21:07:53Z
author_association: MEMBER
diff_hunk: @@ -1,66 +1,138 @@
 import functools
 import operator
+from collections.abc import Callable
+from dataclasses import dataclass
+from types import ModuleType
 
 import numpy as np
 from scipy._lib._array_api import (
     array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray,
-    xp_promote, SCIPY_ARRAY_API
+    xp_promote, xp_capabilities, SCIPY_ARRAY_API
 )
 import scipy._lib.array_api_extra as xpx
 from . import _ufuncs
-# These don't really need to be imported, but otherwise IDEs might not realize
-# that these are defined in this file / report an error in __init__.py
-from ._ufuncs import (
-    log_ndtr, ndtr, ndtri, erf, erfc, i0, i0e, i1, i1e, gammaln,  # noqa: F401
-    gammainc, gammaincc, logit, expit, entr, rel_entr, xlogy,  # noqa: F401
-    chdtr, chdtrc, betainc, betaincc, stdtr, stdtrit  # noqa: F401
-)
-
-array_api_compat_prefix = "scipy._lib.array_api_compat"
 
 
-def get_array_special_func(f_name, xp):
-    if is_numpy(xp):
-        return getattr(_ufuncs, f_name)
-
-    spx = scipy_namespace_for(xp)
-    if spx is not None:
-        f = getattr(spx.special, f_name, None)
+@dataclass
+class _FuncInfo:
+    # NumPy-only function. IT MUST BE ELEMENTWISE.
+    func: Callable
+    # Number of arguments, not counting out=
+    # This is for testing purposes only, due to the fact that
+    # inspect.signature() just returns *args for Cython ufuncs.
+    n_args: int
+    # @xp_capabilities decorator, for the purpose of
+    # documentation and unit testing. Omit to indicate
+    # full support for all backends.
+    xp_capabilities: Callable[[Callable], Callable] | None = None
+    # Generic implementation to fall back on if there is no native dispatch
+    # available. This is a function that accepts (main namespace, scipy namespace)
+    # and returns the final callable, or None if not available.
+    generic_impl: Callable[
+        [ModuleType, ModuleType | None], Callable | None
+    ] | None = None
+
+    @property
+    def name(self):
+        return self.func.__name__
+
+    # These are needed by @lru_cache below
+    def __hash__(self):
+        return hash(self.func)
+
+    def __eq__(self, other):
+        return isinstance(other, _FuncInfo) and self.func == other.func
+
+    @property
+    def wrapper(self):
+        if self.name in globals():
+            # Already initialised. We are likely in a unit test.
+            # Return function potentially overridden by xpx.testing.lazy_xp_function.
+            import scipy.special
+            return getattr(scipy.special, self.name)
+
+        if SCIPY_ARRAY_API:
+            @functools.wraps(self.func)
+            def wrapped(*args, **kwargs):
+                xp = array_namespace(*args)
+                return self._wrapper_for(xp)(*args, **kwargs)
+
+            # Allow pickling the function. Normally this is done by @wraps,
+            # but in this case it doesn't work because self.func is a ufunc.
+            wrapped.__module__ = "scipy.special"
+            wrapped.__qualname__ = self.name
+            func = wrapped
+        else:
+            func = self.func
+
+        capabilities = self.xp_capabilities or xp_capabilities()
+        # In order to retain a naked Cython ufunc when SCIPY_ARRAY_API is
path: scipy/special/_support_alternative_backends.py
user: login: steppi
body: ```suggestion
        # In order to retain a naked ufunc when SCIPY_ARRAY_API is
```
created_at: 2025-05-14T21:08:36Z
author_association: MEMBER
diff_hunk: @@ -93,19 +165,13 @@ def __xlogy(x, y, *, xp=xp):
     return __xlogy
 
 
-def _get_native_func(xp, spx, f_name):
-    f = getattr(spx.special, f_name, None) if spx else None
-    if f is None and hasattr(xp, 'special'):
-        f = getattr(xp.special, f_name, None)
path: scipy/special/_support_alternative_backends.py
user: login: steppi
body: I think these lines were in anticipation of the creation of a `special` Array-API extension, https://github.com/data-apis/array-api/issues/725. This is still planned, but was put on hold at the recommendation of @izaid due to scarcity of available functions across backends, and the hope that his ideas which became [xsf](https://github.com/data-apis/array-api/issues/725) could rectify that situation. I would classify it as defensive programming to include them, and would recommend leaving them in, but with a short comment explaining why they are there.

```python
# In anticipation of special array-api extension
# https://github.com/data-apis/array-api/issues/725
```
created_at: 2025-05-15T01:49:00Z
author_association: MEMBER
diff_hunk: @@ -179,62 +246,45 @@ def fun(t, df, p):  return stdtr(df, t) - p
     return __stdtrit
 
 
-_generic_implementations = {'rel_entr': _rel_entr,
-                            'xlogy': _xlogy,
-                            'chdtr': _chdtr,
-                            'chdtrc': _chdtrc,
-                            'betaincc': _betaincc,
-                            'stdtr': _stdtr,
-                            'stdtrit': _stdtrit,
-                            }
-
-
-# functools.wraps doesn't work because:
-# 'numpy.ufunc' object has no attribute '__module__'
-def support_alternative_backends(f_name):
-    func = getattr(_ufuncs, f_name)
-
-    @functools.wraps(func)
-    def wrapped(*args, **kwargs):
-        xp = array_namespace(*args)
-        f = get_array_special_func(f_name, xp)
-        return f(*args, **kwargs)
-
-    return wrapped
-
-
-# function name: number of args (for testing purposes)
-array_special_func_map = {
-    'log_ndtr': 1,
-    'ndtr': 1,
-    'ndtri': 1,
-    'erf': 1,
-    'erfc': 1,
-    'i0': 1,
-    'i0e': 1,
-    'i1': 1,
-    'i1e': 1,
-    'gammaln': 1,
-    'gammainc': 2,
-    'gammaincc': 2,
-    'logit': 1,
-    'expit': 1,
-    'entr': 1,
-    'rel_entr': 2,
-    'xlogy': 2,
-    'chdtr': 2,
-    'chdtrc': 2,
-    'betainc': 3,
-    'betaincc': 3,
-    'stdtr': 2,
-    'stdtrit': 2,
-}
-
-globals().update(
-    {f_name: support_alternative_backends(f_name)
-     if SCIPY_ARRAY_API
-     else getattr(_ufuncs, f_name)
-     for f_name in array_special_func_map}
+# Inventory of automatically dispatched functions
+# IMPORTANT: these must all be **elementwise** functions!
+
+# PyTorch doesn't implement `betainc`.
+# On torch CPU we can fall back to Cython, but on GPU it won't work.
path: scipy/special/_support_alternative_backends.py
user: login: steppi
body: ```suggestion
# On torch CPU we can fall back to NumPy, but on GPU it won't work.
```
created_at: 2025-05-15T02:06:47Z
author_association: MEMBER
diff_hunk: @@ -179,62 +246,45 @@ def fun(t, df, p):  return stdtr(df, t) - p
     return __stdtrit
 
 
-_generic_implementations = {'rel_entr': _rel_entr,
-                            'xlogy': _xlogy,
-                            'chdtr': _chdtr,
-                            'chdtrc': _chdtrc,
-                            'betaincc': _betaincc,
-                            'stdtr': _stdtr,
-                            'stdtrit': _stdtrit,
-                            }
-
-
-# functools.wraps doesn't work because:
-# 'numpy.ufunc' object has no attribute '__module__'
-def support_alternative_backends(f_name):
-    func = getattr(_ufuncs, f_name)
-
-    @functools.wraps(func)
-    def wrapped(*args, **kwargs):
-        xp = array_namespace(*args)
-        f = get_array_special_func(f_name, xp)
-        return f(*args, **kwargs)
-
-    return wrapped
-
-
-# function name: number of args (for testing purposes)
-array_special_func_map = {
-    'log_ndtr': 1,
-    'ndtr': 1,
-    'ndtri': 1,
-    'erf': 1,
-    'erfc': 1,
-    'i0': 1,
-    'i0e': 1,
-    'i1': 1,
-    'i1e': 1,
-    'gammaln': 1,
-    'gammainc': 2,
-    'gammaincc': 2,
-    'logit': 1,
-    'expit': 1,
-    'entr': 1,
-    'rel_entr': 2,
-    'xlogy': 2,
-    'chdtr': 2,
-    'chdtrc': 2,
-    'betainc': 3,
-    'betaincc': 3,
-    'stdtr': 2,
-    'stdtrit': 2,
-}
-
-globals().update(
-    {f_name: support_alternative_backends(f_name)
-     if SCIPY_ARRAY_API
-     else getattr(_ufuncs, f_name)
-     for f_name in array_special_func_map}
+# Inventory of automatically dispatched functions
+# IMPORTANT: these must all be **elementwise** functions!
+
+# PyTorch doesn't implement `betainc`.
+# On torch CPU we can fall back to Cython, but on GPU it won't work.
+_needs_betainc = xp_capabilities(cpu_only=True, exceptions=['jax.numpy', 'cupy'])
+
+_special_funcs = (
+    _FuncInfo(_ufuncs.betainc, 3, _needs_betainc),
+    _FuncInfo(_ufuncs.betaincc, 3, _needs_betainc, generic_impl=_betaincc),
+    _FuncInfo(_ufuncs.chdtr, 2, generic_impl=_chdtr),
+    _FuncInfo(_ufuncs.chdtrc, 2, generic_impl=_chdtrc),
+    _FuncInfo(_ufuncs.erf, 1),
+    _FuncInfo(_ufuncs.erfc, 1),
+    _FuncInfo(_ufuncs.entr, 1),
+    _FuncInfo(_ufuncs.expit, 1),
+    _FuncInfo(_ufuncs.i0, 1),
+    _FuncInfo(_ufuncs.i0e, 1),
+    _FuncInfo(_ufuncs.i1, 1),
+    _FuncInfo(_ufuncs.i1e, 1),
+    _FuncInfo(_ufuncs.log_ndtr, 1),
+    _FuncInfo(_ufuncs.logit, 1),
+    _FuncInfo(_ufuncs.gammaln, 1),
+    _FuncInfo(_ufuncs.gammainc, 2),
+    _FuncInfo(_ufuncs.gammaincc, 2),
+    _FuncInfo(_ufuncs.ndtr, 1),
+    _FuncInfo(_ufuncs.ndtri, 1),
+    _FuncInfo(_ufuncs.rel_entr, 2, generic_impl=_rel_entr),
+    _FuncInfo(_ufuncs.stdtr,  2, _needs_betainc, generic_impl=_stdtr),
+    _FuncInfo(_ufuncs.stdtrit, 2,
+              xp_capabilities(
+                  cpu_only=True, exceptions=['cupy'],  # needs betainc
+                  skip_backends=[("jax.numpy", "no scipy.optimize support")]),
+              generic_impl=_stdtrit),
+    _FuncInfo(_ufuncs.xlogy, 2, generic_impl=_xlogy),
 )
 
-__all__ = list(array_special_func_map)
+# Override ufuncs.
+# When SCIPY_ARRAY_API is disabled, this exclusively updates the docstrings in place
+# and populates the xp_capabilities table, while retaining the original Cython ufuncs.
path: scipy/special/_support_alternative_backends.py
user: login: steppi
body: ```suggestion
# and populates the xp_capabilities table, while retaining the original ufuncs.
```
created_at: 2025-05-15T02:07:13Z
author_association: MEMBER
diff_hunk: @@ -197,21 +172,70 @@ def test_support_alternative_backends_hypothesis(xp, f_name, n_args, data):
     if positive_only:
         elements['min_value'] = 0
 
-    shapes, _ = data.draw(npst.mutually_broadcastable_shapes(num_shapes=n_args))
+    shapes, _ = data.draw(
+        npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args))
     args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements))
                for shape in shapes]
 
     args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np]
     args_np = [np.asarray(arg, dtype=dtype_np_ref) for arg in args_np]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped Cython ufunc
 
     # When dtype_np is integer, the output dtype can be float
     atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps
     xp_assert_close(res, xp.asarray(ref), atol=atol)
 
 
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_pickle(func):
+    roundtrip = pickle.loads(pickle.dumps(func))
+    assert roundtrip is func
+
+
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_repr(func):
+    assert func.__name__ in repr(func)
+    assert "locals" not in repr(func)
+
+
+@pytest.mark.skipif(
+    np.__version__ < "2",
path: scipy/special/tests/test_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    version.parse(np.__version__) < version.parse("2.2"),
```

and need to add the import

```python
from packaging import version
```
created_at: 2025-05-15T14:50:17Z
author_association: MEMBER
diff_hunk: @@ -197,21 +172,70 @@ def test_support_alternative_backends_hypothesis(xp, f_name, n_args, data):
     if positive_only:
         elements['min_value'] = 0
 
-    shapes, _ = data.draw(npst.mutually_broadcastable_shapes(num_shapes=n_args))
+    shapes, _ = data.draw(
+        npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args))
     args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements))
                for shape in shapes]
 
     args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np]
     args_np = [np.asarray(arg, dtype=dtype_np_ref) for arg in args_np]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped Cython ufunc
 
     # When dtype_np is integer, the output dtype can be float
     atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps
     xp_assert_close(res, xp.asarray(ref), atol=atol)
 
 
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_pickle(func):
+    roundtrip = pickle.loads(pickle.dumps(func))
+    assert roundtrip is func
+
+
+@pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs])
+def test_repr(func):
+    assert func.__name__ in repr(func)
+    assert "locals" not in repr(func)
+
+
+@pytest.mark.skipif(
+    version.parse(np.__version__) < version.parse("2.2"),
+    reason="Can't update ufunc __doc__ when SciPy is compiled vs. NumPy 1.x")
path: scipy/special/tests/test_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    reason="Can't update ufunc __doc__ when SciPy is compiled vs. NumPy < 2.2")
```
created_at: 2025-05-15T15:06:13Z
author_association: MEMBER
diff_hunk: @@ -127,27 +105,24 @@ def test_support_alternative_backends(xp, f_name, n_args, dtype, shapes):
         # Try to trigger bugs related to having multiple chunks.
         args_xp = [arg.rechunk(5) for arg in args_xp]
 
-    res = fw(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped Cython ufunc
path: scipy/special/tests/test_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    ref = nfo.func(*args_np)  # Unwrapped ufunc
```
created_at: 2025-05-15T15:06:40Z
author_association: MEMBER
diff_hunk: @@ -165,8 +140,8 @@ def test_support_alternative_backends_mismatched_dtypes(xp, f_name, n_args):
     args_np = [np.asarray(arg, dtype=dtype_np_ref)
                for arg, dtype_np_ref in zip(args_np, dtypes_np_ref)]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped Cython ufunc
path: scipy/special/tests/test_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    ref = nfo.func(*args_np)  # Unwrapped ufunc
```
created_at: 2025-05-15T15:07:00Z
author_association: MEMBER
diff_hunk: @@ -197,21 +172,70 @@ def test_support_alternative_backends_hypothesis(xp, f_name, n_args, data):
     if positive_only:
         elements['min_value'] = 0
 
-    shapes, _ = data.draw(npst.mutually_broadcastable_shapes(num_shapes=n_args))
+    shapes, _ = data.draw(
+        npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args))
     args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements))
                for shape in shapes]
 
     args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np]
     args_np = [np.asarray(arg, dtype=dtype_np_ref) for arg in args_np]
 
-    res = f(*args_xp)
-    ref = f(*args_np)
+    res = nfo.wrapper(*args_xp)  # Also wrapped by lazy_xp_function
+    ref = nfo.func(*args_np)  # Unwrapped Cython ufunc
path: scipy/special/tests/test_support_alternative_backends.py
user: login: steppi
body: ```suggestion
    ref = nfo.func(*args_np)  # Unwrapped ufunc
```
created_at: 2025-05-15T15:07:17Z
author_association: MEMBER
diff_hunk: @@ -93,19 +165,13 @@ def __xlogy(x, y, *, xp=xp):
     return __xlogy
 
 
-def _get_native_func(xp, spx, f_name):
-    f = getattr(spx.special, f_name, None) if spx else None
-    if f is None and hasattr(xp, 'special'):
-        f = getattr(xp.special, f_name, None)
path: scipy/special/_support_alternative_backends.py
user: login: crusaderky
body: reverted
created_at: 2025-05-15T15:11:10Z
author_association: CONTRIBUTOR

PR comments:
user: login: crusaderky
created_at: 2025-05-15T15:13:54Z
body: @steppi all comments have been addressed
author_association: CONTRIBUTOR
user: login: crusaderky
created_at: 2025-05-15T15:54:37Z
body: windows failures are unrelated
author_association: CONTRIBUTOR
user: login: steppi
created_at: 2025-05-15T18:04:23Z
body: Cool, the free-threaded failure is unrelated. I'll go ahead and merge this. @tylerjereddy, can this make it into 1.16?
author_association: MEMBER
user: login: crusaderky
created_at: 2025-05-15T18:05:08Z
body: free-threaded CI failed somewhere in `scipy/stats/_distn_infrastructure.py` and I'm inclined to think it's unrelated.
Unless there are known issues with `functools.lru_cache`? @rgommers 
author_association: CONTRIBUTOR
user: login: tylerjereddy
created_at: 2025-05-15T20:19:27Z
body: >  I'll go ahead and merge this. @tylerjereddy, can this make it into 1.16?

You merged it so I don't have much choice now unless I revert it. I have no plans to revert it unless it causes problems for me during the release process though.
author_association: CONTRIBUTOR
user: login: rgommers
created_at: 2025-05-16T03:10:37Z
body: > free-threaded CI failed somewhere in `scipy/stats/_distn_infrastructure.py` and I'm inclined to think it's unrelated.
> Unless there are known issues with `functools.lru_cache`? @rgommers

It's unrelated to this PR, I just added these segfaults to https://github.com/scipy/scipy/issues/22758#issuecomment-2883968653, they were recently introduced in a `stats` PR.
author_association: MEMBER
user: login: rgommers
created_at: 2025-05-16T03:12:04Z
body: Thanks for reviewing this @steppi, great to see this get merged in time for the next release. 
author_association: MEMBER