MAINT: _lib._array_api: clean-up

Author: lucascolley

scipy/_lib/_array_api.py

@@ -20,12 +20,21 @@
 from scipy._lib import array_api_compat
 from scipy._lib.array_api_compat import (
     is_array_api_obj,
-    size,
+    size as xp_size,
     numpy as np_compat,
-    device
+    device as xp_device
 )
 
-__all__ = ['array_namespace', '_asarray', 'size', 'device']
+__all__ = [
+    '_asarray', 'array_namespace', 'assert_almost_equal', 'assert_array_almost_equal',
+    'get_xp_devices',
+    '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_moveaxis_to_end', 'xp_ravel', 'xp_real', 'xp_sign', 'xp_size',
+    'xp_take_along_axis', 'xp_unsupported_param_msg', 'xp_vector_norm',
+]
 
 
 # To enable array API and strict array-like input validation
@@ -44,7 +53,7 @@
     ArrayLike = Array | npt.ArrayLike
 
 
-def compliance_scipy(arrays: list[ArrayLike]) -> list[Array]:
+def _compliance_scipy(arrays: list[ArrayLike]) -> list[Array]:
     """Raise exceptions on known-bad subclasses.
 
     The following subclasses are not supported and raise and error:
@@ -111,7 +120,7 @@ def array_namespace(*arrays: Array) -> ModuleType:
 
     1. Check for the global switch: SCIPY_ARRAY_API. This can also be accessed
        dynamically through ``_GLOBAL_CONFIG['SCIPY_ARRAY_API']``.
-    2. `compliance_scipy` raise exceptions on known-bad subclasses. See
+    2. `_compliance_scipy` raise exceptions on known-bad subclasses. See
        its definition for more details.
 
     When the global switch is False, it defaults to the `numpy` namespace.
@@ -124,7 +133,7 @@ def array_namespace(*arrays: Array) -> ModuleType:
 
     _arrays = [array for array in arrays if array is not None]
 
-    _arrays = compliance_scipy(_arrays)
+    _arrays = _compliance_scipy(_arrays)
 
     return array_api_compat.array_namespace(*_arrays)
 
@@ -176,18 +185,18 @@ def _asarray(
     return array
 
 
-def atleast_nd(x: Array, *, ndim: int, xp: ModuleType | None = None) -> 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 = atleast_nd(x, ndim=ndim, xp=xp)
+        x = xp_atleast_nd(x, ndim=ndim, xp=xp)
     return x
 
 
-def copy(x: Array, *, xp: ModuleType | None = None) -> Array:
+def xp_copy(x: Array, *, xp: ModuleType | None = None) -> Array:
     """
     Copies an array.
 
@@ -207,7 +216,8 @@ def copy(x: Array, *, xp: ModuleType | None = None) -> Array:
     This copy function does not offer all the semantics of `np.copy`, i.e. the
     `subok` and `order` keywords are not used.
     """
-    # Note: xp.asarray fails if xp is numpy.
+    # Note: for older NumPy versions, `np.asarray` did not support the `copy` kwarg,
+    # so this uses our other helper `_asarray`.
     if xp is None:
         xp = array_namespace(x)
 
@@ -395,14 +405,14 @@ def assert_almost_equal(actual, desired, decimal=7, *args, **kwds):
                            *args, **kwds)
 
 
-def cov(x: Array, *, xp: ModuleType | None = None) -> Array:
+def xp_cov(x: Array, *, xp: ModuleType | None = None) -> Array:
     if xp is None:
         xp = array_namespace(x)
 
-    X = copy(x, xp=xp)
+    X = xp_copy(x, xp=xp)
     dtype = xp.result_type(X, xp.float64)
 
-    X = atleast_nd(X, ndim=2, xp=xp)
+    X = xp_atleast_nd(X, ndim=2, xp=xp)
     X = xp.asarray(X, dtype=dtype)
 
     avg = xp.mean(X, axis=1)

scipy/_lib/_elementwise_iterative_method.py

@@ -14,7 +14,7 @@
 import math
 import numpy as np
 from ._util import _RichResult, _call_callback_maybe_halt
-from ._array_api import array_namespace, size as xp_size
+from ._array_api import array_namespace, xp_size
 
 _ESIGNERR = -1
 _ECONVERR = -2

scipy/_lib/_util.py

@@ -15,7 +15,7 @@
 )
 
 import numpy as np
-from scipy._lib._array_api import array_namespace, is_numpy, size as xp_size
+from scipy._lib._array_api import array_namespace, is_numpy, xp_size
 
 
 AxisError: type[Exception]

scipy/_lib/tests/test__util.py

@@ -13,7 +13,7 @@
 from scipy.conftest import array_api_compatible, skip_xp_invalid_arg
 
 from scipy._lib._array_api import (xp_assert_equal, xp_assert_close, is_numpy,
-                                   copy as xp_copy, is_array_api_strict)
+                                   xp_copy, is_array_api_strict)
 from scipy._lib._util import (_aligned_zeros, check_random_state, MapWrapper,
                               getfullargspec_no_self, FullArgSpec,
                               rng_integers, _validate_int, _rename_parameter,

scipy/_lib/tests/test_array_api.py

@@ -3,7 +3,7 @@
 
 from scipy.conftest import array_api_compatible
 from scipy._lib._array_api import (
-    _GLOBAL_CONFIG, array_namespace, _asarray, copy, xp_assert_equal, is_numpy
+    _GLOBAL_CONFIG, array_namespace, _asarray, xp_copy, xp_assert_equal, is_numpy
 )
 import scipy._lib.array_api_compat.numpy as np_compat
 
@@ -60,7 +60,7 @@ def test_array_likes(self):
     def test_copy(self, xp):
         for _xp in [xp, None]:
             x = xp.asarray([1, 2, 3])
-            y = copy(x, xp=_xp)
+            y = xp_copy(x, xp=_xp)
             # with numpy we'd want to use np.shared_memory, but that's not specified
             # in the array-api
             x[0] = 10

scipy/cluster/hierarchy.py

@@ -134,7 +134,7 @@
 import numpy as np
 from . import _hierarchy, _optimal_leaf_ordering
 import scipy.spatial.distance as distance
-from scipy._lib._array_api import array_namespace, _asarray, copy, is_jax
+from scipy._lib._array_api import array_namespace, _asarray, xp_copy, is_jax
 from scipy._lib._disjoint_set import DisjointSet
 
 
@@ -1358,7 +1358,7 @@ def cut_tree(Z, n_clusters=None, height=None):
 
     for i, node in enumerate(nodes):
         idx = node.pre_order()
-        this_group = copy(last_group, xp=xp)
+        this_group = xp_copy(last_group, xp=xp)
         # TODO ARRAY_API complex indexing not supported
         this_group[idx] = xp.min(last_group[idx])
         this_group[this_group > xp.max(last_group[idx])] -= 1
@@ -1822,14 +1822,14 @@ def from_mlab_linkage(Z):
 
     # If it's empty, return it.
     if len(Zs) == 0 or (len(Zs) == 1 and Zs[0] == 0):
-        return copy(Z, xp=xp)
+        return xp_copy(Z, xp=xp)
 
     if len(Zs) != 2:
         raise ValueError("The linkage array must be rectangular.")
 
     # If it contains no rows, return it.
     if Zs[0] == 0:
-        return copy(Z, xp=xp)
+        return xp_copy(Z, xp=xp)
 
     if xp.min(Z[:, 0:2]) != 1.0 and xp.max(Z[:, 0:2]) != 2 * Zs[0]:
         raise ValueError('The format of the indices is not 1..N')
@@ -1925,7 +1925,7 @@ def to_mlab_linkage(Z):
     Z = _asarray(Z, order='C', dtype=xp.float64, xp=xp)
     Zs = Z.shape
     if len(Zs) == 0 or (len(Zs) == 1 and Zs[0] == 0):
-        return copy(Z, xp=xp)
+        return xp_copy(Z, xp=xp)
     is_valid_linkage(Z, throw=True, name='Z')
 
     return xp.concat((Z[:, :2] + 1.0, Z[:, 2:3]), axis=1)

scipy/cluster/tests/test_vq.py

@@ -16,7 +16,7 @@
 from scipy.sparse._sputils import matrix
 
 from scipy._lib._array_api import (
-    SCIPY_ARRAY_API, copy, cov, xp_assert_close, xp_assert_equal
+    SCIPY_ARRAY_API, xp_copy, xp_cov, xp_assert_close, xp_assert_equal
 )
 
 pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
@@ -307,7 +307,7 @@ def test_kmeans2_rank1(self, xp):
         data1 = data[:, 0]
 
         initc = data1[:3]
-        code = copy(initc, xp=xp)
+        code = xp_copy(initc, xp=xp)
         kmeans2(data1, code, iter=1)[0]
         kmeans2(data1, code, iter=2)[0]
 
@@ -353,8 +353,8 @@ def test_krandinit(self, xp):
         for data in datas:
             rng = np.random.default_rng(1234)
             init = _krandinit(data, k, rng, xp)
-            orig_cov = cov(data.T)
-            init_cov = cov(init.T)
+            orig_cov = xp_cov(data.T)
+            init_cov = xp_cov(init.T)
             xp_assert_close(orig_cov, init_cov, atol=1.1e-2)
 
     def test_kmeans2_empty(self, xp):

scipy/cluster/vq.py

@@ -68,7 +68,7 @@
 import numpy as np
 from collections import deque
 from scipy._lib._array_api import (
-    _asarray, array_namespace, size, atleast_nd, copy, cov
+    _asarray, array_namespace, xp_size, xp_atleast_nd, xp_copy, xp_cov
 )
 from scipy._lib._util import check_random_state, rng_integers
 from scipy.spatial.distance import cdist
@@ -472,8 +472,8 @@ def kmeans(obs, k_or_guess, iter=20, thresh=1e-5, check_finite=True,
         raise ValueError(f"iter must be at least 1, got {iter}")
 
     # Determine whether a count (scalar) or an initial guess (array) was passed.
-    if size(guess) != 1:
-        if size(guess) < 1:
+    if xp_size(guess) != 1:
+        if xp_size(guess) < 1:
             raise ValueError(f"Asked for 0 clusters. Initial book was {guess}")
         return _kmeans(obs, guess, thresh=thresh, xp=xp)
 
@@ -551,23 +551,23 @@ def _krandinit(data, k, rng, xp):
     k = np.asarray(k)
 
     if data.ndim == 1:
-        _cov = cov(data)
+        _cov = xp_cov(data)
         x = rng.standard_normal(size=k)
         x = xp.asarray(x)
         x *= xp.sqrt(_cov)
     elif data.shape[1] > data.shape[0]:
         # initialize when the covariance matrix is rank deficient
         _, s, vh = xp.linalg.svd(data - mu, full_matrices=False)
-        x = rng.standard_normal(size=(k, size(s)))
+        x = rng.standard_normal(size=(k, xp_size(s)))
         x = xp.asarray(x)
         sVh = s[:, None] * vh / xp.sqrt(data.shape[0] - xp.asarray(1.))
         x = x @ sVh
     else:
-        _cov = atleast_nd(cov(data.T), ndim=2)
+        _cov = xp_atleast_nd(xp_cov(data.T), ndim=2)
 
         # k rows, d cols (one row = one obs)
         # Generate k sample of a random variable ~ Gaussian(mu, cov)
-        x = rng.standard_normal(size=(k, size(mu)))
+        x = rng.standard_normal(size=(k, xp_size(mu)))
         x = xp.asarray(x)
         x = x @ xp.linalg.cholesky(_cov).T
 
@@ -782,19 +782,19 @@ def kmeans2(data, k, iter=10, thresh=1e-5, minit='random',
     else:
         xp = array_namespace(data, k)
     data = _asarray(data, xp=xp, check_finite=check_finite)
-    code_book = copy(k, xp=xp)
+    code_book = xp_copy(k, xp=xp)
     if data.ndim == 1:
         d = 1
     elif data.ndim == 2:
         d = data.shape[1]
     else:
         raise ValueError("Input of rank > 2 is not supported.")
 
-    if size(data) < 1 or size(code_book) < 1:
+    if xp_size(data) < 1 or xp_size(code_book) < 1:
         raise ValueError("Empty input is not supported.")
 
     # If k is not a single value, it should be compatible with data's shape
-    if minit == 'matrix' or size(code_book) > 1:
+    if minit == 'matrix' or xp_size(code_book) > 1:
         if data.ndim != code_book.ndim:
             raise ValueError("k array doesn't match data rank")
         nc = code_book.shape[0]

scipy/fft/_fftlog_backend.py

@@ -3,7 +3,7 @@
 from ._basic import rfft, irfft
 from ..special import loggamma, poch
 
-from scipy._lib._array_api import array_namespace, copy
+from scipy._lib._array_api import array_namespace
 
 __all__ = ['fht', 'ifht', 'fhtoffset']
 
@@ -106,12 +106,12 @@ def fhtcoeff(n, dln, mu, offset=0.0, bias=0.0, inverse=False):
     if np.isinf(u[0]) and not inverse:
         warn('singular transform; consider changing the bias', stacklevel=3)
         # fix coefficient to obtain (potentially correct) transform anyway
-        u = copy(u)
+        u = np.copy(u)
         u[0] = 0
     elif u[0] == 0 and inverse:
         warn('singular inverse transform; consider changing the bias', stacklevel=3)
         # fix coefficient to obtain (potentially correct) inverse anyway
-        u = copy(u)
+        u = np.copy(u)
         u[0] = np.inf
 
     return u

scipy/fft/tests/test_basic.py

@@ -9,7 +9,7 @@
 import scipy.fft as fft
 from scipy.conftest import array_api_compatible
 from scipy._lib._array_api import (
-    array_namespace, size, xp_assert_close, xp_assert_equal
+    array_namespace, xp_size, xp_assert_close, xp_assert_equal
 )
 
 pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
@@ -123,7 +123,7 @@ def test_ifftn(self, xp):
 
     def test_rfft(self, xp):
         x = xp.asarray(random(29), dtype=xp.float64)
-        for n in [size(x), 2*size(x)]:
+        for n in [xp_size(x), 2*xp_size(x)]:
             for norm in [None, "backward", "ortho", "forward"]:
                 xp_assert_close(fft.rfft(x, n=n, norm=norm),
                                 fft.fft(xp.asarray(x, dtype=xp.complex128),
@@ -285,7 +285,7 @@ def test_all_1d_norm_preserving(self, xp):
         x = xp.asarray(random(30), dtype=xp.float64)
         xp_test = array_namespace(x)
         x_norm = xp_test.linalg.vector_norm(x)
-        n = size(x) * 2
+        n = xp_size(x) * 2
         func_pairs = [(fft.rfft, fft.irfft),
                       # hfft: order so the first function takes x.size samples
                       #       (necessary for comparison to x_norm above)
@@ -297,7 +297,7 @@ def test_all_1d_norm_preserving(self, xp):
             if forw == fft.fft:
                 x = xp.asarray(x, dtype=xp.complex128)
                 x_norm = xp_test.linalg.vector_norm(x)
-            for n in [size(x), 2*size(x)]:
+            for n in [xp_size(x), 2*xp_size(x)]:
                 for norm in ['backward', 'ortho', 'forward']:
                     tmp = forw(x, n=n, norm=norm)
                     tmp = back(tmp, n=n, norm=norm)