MAINT: _lib: vendor and use array-api-extra (scipy#21600)

Author: lucascolley

.gitmodules

@@ -26,3 +26,6 @@
 [submodule "scipy/_lib/cobyqa"]
 	path = scipy/_lib/cobyqa
 	url = https://github.com/cobyqa/cobyqa.git
+[submodule "scipy/_lib/array_api_extra"]
+	path = scipy/_lib/array_api_extra
+	url = https://github.com/lucascolley/array-api-extra.git

scipy/_lib/_array_api.py

@@ -9,7 +9,6 @@
 from __future__ import annotations
 
 import os
-import warnings
 
 from types import ModuleType
 from typing import Any, Literal, TYPE_CHECKING
@@ -36,7 +35,7 @@
     'is_array_api_strict', 'is_complex', 'is_cupy', 'is_jax', 'is_numpy', 'is_torch', 
     'SCIPY_ARRAY_API', 'SCIPY_DEVICE', 'scipy_namespace_for',
     'xp_assert_close', 'xp_assert_equal', 'xp_assert_less',
-    'xp_atleast_nd', 'xp_copy', 'xp_copysign', 'xp_device',
+    'xp_copy', 'xp_copysign', 'xp_device',
     'xp_moveaxis_to_end', 'xp_ravel', 'xp_real', 'xp_sign', 'xp_size',
     'xp_take_along_axis', 'xp_unsupported_param_msg', 'xp_vector_norm',
     'xp_create_diagonal'
@@ -200,17 +199,6 @@ def _asarray(
     return array
 
 
-def xp_atleast_nd(x: Array, *, ndim: int, xp: ModuleType | None = None) -> Array:
-    """Recursively expand the dimension to have at least `ndim`."""
-    if xp is None:
-        xp = array_namespace(x)
-    x = xp.asarray(x)
-    if x.ndim < ndim:
-        x = xp.expand_dims(x, axis=0)
-        x = xp_atleast_nd(x, ndim=ndim, xp=xp)
-    return x
-
-
 def xp_copy(x: Array, *, xp: ModuleType | None = None) -> Array:
     """
     Copies an array.
@@ -382,34 +370,6 @@ def assert_almost_equal(actual, desired, decimal=7, *args, **kwds):
                            *args, **kwds)
 
 
-def xp_cov(x: Array, *, xp: ModuleType | None = None) -> Array:
-    if xp is None:
-        xp = array_namespace(x)
-
-    X = xp_copy(x, xp=xp)
-    dtype = xp.result_type(X, xp.float64)
-
-    X = xp_atleast_nd(X, ndim=2, xp=xp)
-    X = xp.asarray(X, dtype=dtype)
-
-    avg = xp.mean(X, axis=1)
-    fact = X.shape[1] - 1
-
-    if fact <= 0:
-        warnings.warn("Degrees of freedom <= 0 for slice",
-                      RuntimeWarning, stacklevel=2)
-        fact = 0.0
-
-    X -= avg[:, None]
-    X_T = X.T
-    if xp.isdtype(X_T.dtype, 'complex floating'):
-        X_T = xp.conj(X_T)
-    c = X @ X_T
-    c /= fact
-    axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
-    return xp.squeeze(c, axis=axes)
-
-
 def xp_unsupported_param_msg(param: Any) -> str:
     return f'Providing {param!r} is only supported for numpy arrays.'
 

scipy/_lib/array_api_extra

@@ -11,6 +11,9 @@ endif
 if not fs.exists('array_api_compat/README.md')
   error('Missing the `array_api_compat` submodule! Run `git submodule update --init` to fix this.')
 endif
+if not fs.exists('array_api_extra/README.md')
+  error('Missing the `array_api_extra` submodule! Run `git submodule update --init` to fix this.')
+endif
 if not fs.exists('pocketfft/README.md')
   error('Missing the `pocketfft` submodule! Run `git submodule update --init` to fix this.')
 endif
@@ -213,6 +216,18 @@ py3.install_sources(
   subdir: 'scipy/_lib/array_api_compat/torch',
 )
 
+# `array_api_extra` install to simplify import path;
+# should be updated whenever new files are added to `array_api_extra`
+
+py3.install_sources(
+  [
+    'array_api_extra/src/array_api_extra/__init__.py',
+    'array_api_extra/src/array_api_extra/_funcs.py',
+    'array_api_extra/src/array_api_extra/_typing.py',
+  ],
+  subdir: 'scipy/_lib/array_api_extra',
+)
+
 py3.install_sources(
   [
     'cobyqa/cobyqa/__init__.py',

scipy/_lib/meson.build

@@ -15,8 +15,9 @@
 from scipy.conftest import array_api_compatible
 from scipy.sparse._sputils import matrix
 
+from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api import (
-    SCIPY_ARRAY_API, xp_copy, xp_cov, xp_assert_close, xp_assert_equal
+    SCIPY_ARRAY_API, array_namespace, xp_copy, xp_assert_close, xp_assert_equal
 )
 
 pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
@@ -352,11 +353,12 @@ def test_krandinit(self, xp):
         datas = [xp.reshape(data, (200, 2)),
                  xp.reshape(data, (20, 20))[:10, :]]
         k = int(1e6)
+        xp_test = array_namespace(data)
         for data in datas:
             rng = np.random.default_rng(1234)
-            init = _krandinit(data, k, rng, xp)
-            orig_cov = xp_cov(data.T)
-            init_cov = xp_cov(init.T)
+            init = _krandinit(data, k, rng, xp_test)
+            orig_cov = xpx.cov(data.T, xp=xp_test)
+            init_cov = xpx.cov(init.T, xp=xp_test)
             xp_assert_close(orig_cov, init_cov, atol=1.1e-2)
 
     def test_kmeans2_empty(self, xp):

scipy/cluster/tests/test_vq.py

@@ -68,10 +68,11 @@
 import numpy as np
 from collections import deque
 from scipy._lib._array_api import (
-    _asarray, array_namespace, xp_size, xp_atleast_nd, xp_copy, xp_cov
+    _asarray, array_namespace, xp_size, xp_copy
 )
 from scipy._lib._util import (check_random_state, rng_integers,
                               _transition_to_rng)
+from scipy._lib import array_api_extra as xpx
 from scipy.spatial.distance import cdist
 
 from . import _vq
@@ -548,7 +549,7 @@ def _krandinit(data, k, rng, xp):
     k = np.asarray(k)
 
     if data.ndim == 1:
-        _cov = xp_cov(data)
+        _cov = xpx.cov(data, xp=xp)
         x = rng.standard_normal(size=k)
         x = xp.asarray(x)
         x *= xp.sqrt(_cov)
@@ -560,7 +561,7 @@ def _krandinit(data, k, rng, xp):
         sVh = s[:, None] * vh / xp.sqrt(data.shape[0] - xp.asarray(1.))
         x = x @ sVh
     else:
-        _cov = xp_atleast_nd(xp_cov(data.T), ndim=2)
+        _cov = xpx.atleast_nd(xpx.cov(data.T, xp=xp), ndim=2, xp=xp)
 
         # k rows, d cols (one row = one obs)
         # Generate k sample of a random variable ~ Gaussian(mu, cov)

scipy/cluster/vq.py

@@ -3,7 +3,8 @@
 from ._numdiff import approx_derivative, group_columns
 from ._hessian_update_strategy import HessianUpdateStrategy
 from scipy.sparse.linalg import LinearOperator
-from scipy._lib._array_api import xp_atleast_nd, array_namespace
+from scipy._lib._array_api import array_namespace
+from scipy._lib import array_api_extra as xpx
 
 
 FD_METHODS = ('2-point', '3-point', 'cs')
@@ -183,7 +184,7 @@ def __init__(self, fun, x0, args, grad, hess, finite_diff_rel_step,
                              "quasi-Newton strategies.")
 
         self.xp = xp = array_namespace(x0)
-        _x = xp_atleast_nd(x0, ndim=1, xp=xp)
+        _x = xpx.atleast_nd(xp.asarray(x0), ndim=1, xp=xp)
         _dtype = xp.float64
         if xp.isdtype(_x.dtype, "real floating"):
             _dtype = _x.dtype
@@ -274,7 +275,7 @@ def _update_x(self, x):
             # ensure that self.x is a copy of x. Don't store a reference
             # otherwise the memoization doesn't work properly.
 
-            _x = xp_atleast_nd(x, ndim=1, xp=self.xp)
+            _x = xpx.atleast_nd(self.xp.asarray(x), ndim=1, xp=self.xp)
             self.x = self.xp.astype(_x, self.x_dtype)
             self.f_updated = False
             self.g_updated = False
@@ -283,7 +284,7 @@ def _update_x(self, x):
         else:
             # ensure that self.x is a copy of x. Don't store a reference
             # otherwise the memoization doesn't work properly.
-            _x = xp_atleast_nd(x, ndim=1, xp=self.xp)
+            _x = xpx.atleast_nd(self.xp.asarray(x), ndim=1, xp=self.xp)
             self.x = self.xp.astype(_x, self.x_dtype)
             self.f_updated = False
             self.g_updated = False
@@ -380,7 +381,7 @@ def __init__(self, fun, x0, jac, hess,
                              "strategies.")
 
         self.xp = xp = array_namespace(x0)
-        _x = xp_atleast_nd(x0, ndim=1, xp=xp)
+        _x = xpx.atleast_nd(xp.asarray(x0), ndim=1, xp=xp)
         _dtype = xp.float64
         if xp.isdtype(_x.dtype, "real floating"):
             _dtype = _x.dtype
@@ -557,15 +558,15 @@ def update_x(x):
                 self._update_jac()
                 self.x_prev = self.x
                 self.J_prev = self.J
-                _x = xp_atleast_nd(x, ndim=1, xp=self.xp)
+                _x = xpx.atleast_nd(self.xp.asarray(x), ndim=1, xp=self.xp)
                 self.x = self.xp.astype(_x, self.x_dtype)
                 self.f_updated = False
                 self.J_updated = False
                 self.H_updated = False
                 self._update_hess()
         else:
             def update_x(x):
-                _x = xp_atleast_nd(x, ndim=1, xp=self.xp)
+                _x = xpx.atleast_nd(self.xp.asarray(x), ndim=1, xp=self.xp)
                 self.x = self.xp.astype(_x, self.x_dtype)
                 self.f_updated = False
                 self.J_updated = False
@@ -637,7 +638,7 @@ def __init__(self, A, x0, sparse_jacobian):
         self.m, self.n = self.J.shape
 
         self.xp = xp = array_namespace(x0)
-        _x = xp_atleast_nd(x0, ndim=1, xp=xp)
+        _x = xpx.atleast_nd(xp.asarray(x0), ndim=1, xp=xp)
         _dtype = xp.float64
         if xp.isdtype(_x.dtype, "real floating"):
             _dtype = _x.dtype
@@ -654,7 +655,7 @@ def __init__(self, A, x0, sparse_jacobian):
 
     def _update_x(self, x):
         if not np.array_equal(x, self.x):
-            _x = xp_atleast_nd(x, ndim=1, xp=self.xp)
+            _x = xpx.atleast_nd(self.xp.asarray(x), ndim=1, xp=self.xp)
             self.x = self.xp.astype(_x, self.x_dtype)
             self.f_updated = False
 

scipy/optimize/_differentiable_functions.py

@@ -6,7 +6,8 @@
 from scipy.sparse.linalg import LinearOperator
 from ..sparse import issparse, csc_matrix, csr_matrix, coo_matrix, find
 from ._group_columns import group_dense, group_sparse
-from scipy._lib._array_api import xp_atleast_nd, array_namespace
+from scipy._lib._array_api import array_namespace
+from scipy._lib import array_api_extra as xpx
 
 
 def _adjust_scheme_to_bounds(x0, h, num_steps, scheme, lb, ub):
@@ -440,7 +441,7 @@ def approx_derivative(fun, x0, method='3-point', rel_step=None, abs_step=None,
         raise ValueError(f"Unknown method '{method}'. ")
 
     xp = array_namespace(x0)
-    _x = xp_atleast_nd(x0, ndim=1, xp=xp)
+    _x = xpx.atleast_nd(xp.asarray(x0), ndim=1, xp=xp)
     _dtype = xp.float64
     if xp.isdtype(_x.dtype, "real floating"):
         _dtype = _x.dtype

scipy/optimize/_numdiff.py

@@ -41,8 +41,8 @@
 from scipy._lib._util import (MapWrapper, check_random_state, _RichResult,
                               _call_callback_maybe_halt, _transition_to_rng)
 from scipy.optimize._differentiable_functions import ScalarFunction, FD_METHODS
-from scipy._lib._array_api import (array_namespace, xp_atleast_nd,
-                                   xp_create_diagonal)
+from scipy._lib._array_api import array_namespace, xp_create_diagonal
+from scipy._lib import array_api_extra as xpx
 
 
 # standard status messages of optimizers
@@ -442,7 +442,7 @@ def rosen_hess(x):
 
     """
     xp = array_namespace(x)
-    x = xp_atleast_nd(x, ndim=1, xp=xp)
+    x = xpx.atleast_nd(x, ndim=1, xp=xp)
     if xp.isdtype(x.dtype, 'integral'):
         x = xp.astype(x, xp.asarray(1.).dtype)
     H = (xp_create_diagonal(-400 * x[:-1], offset=1, xp=xp) 
@@ -486,7 +486,7 @@ def rosen_hess_prod(x, p):
 
     """
     xp = array_namespace(x, p)
-    x = xp_atleast_nd(x, ndim=1, xp=xp)
+    x = xpx.atleast_nd(x, ndim=1, xp=xp)
     if xp.isdtype(x.dtype, 'integral'):
         x = xp.astype(x, xp.asarray(1.).dtype)
     p = xp.asarray(p, dtype=x.dtype)

scipy/optimize/_optimize.py

@@ -24,7 +24,8 @@
                         _check_clip_x)
 from ._numdiff import approx_derivative
 from ._constraints import old_bound_to_new, _arr_to_scalar
-from scipy._lib._array_api import xp_atleast_nd, array_namespace
+from scipy._lib._array_api import array_namespace
+from scipy._lib import array_api_extra as xpx
 
 
 __docformat__ = "restructuredtext en"
@@ -250,7 +251,7 @@ def _minimize_slsqp(func, x0, args=(), jac=None, bounds=None,
 
     # Transform x0 into an array.
     xp = array_namespace(x0)
-    x0 = xp_atleast_nd(x0, ndim=1, xp=xp)
+    x0 = xpx.atleast_nd(xp.asarray(x0), ndim=1, xp=xp)
     dtype = xp.float64
     if xp.isdtype(x0.dtype, "real floating"):
         dtype = x0.dtype