TST: array types: wrap namespaces centrally (scipy#22320)

Author: crusaderky

scipy/_lib/_array_api.py

@@ -252,9 +252,6 @@ def _strict_check(actual, desired, xp, *,
             xp = _default_xp_ctxvar.get()
         except LookupError:
             xp = array_namespace(desired)
-        else:
-            # Wrap namespace if needed
-            xp = array_namespace(xp.asarray(0))
 
     if check_namespace:
         _assert_matching_namespace(actual, desired, xp)

scipy/_lib/tests/test__util.py

@@ -615,8 +615,8 @@ def f2(*args):
 
         # Ensure arrays are at least 1d to follow sane type promotion rules.
         # This can be removed when minimum supported NumPy is 2.0
-        if xp == np:
-            cond, fillvalue, *arrays = np.atleast_1d(cond, fillvalue, *arrays)
+        if is_numpy(xp):
+            cond, fillvalue, *arrays = xp.atleast_1d(cond, fillvalue, *arrays)
 
         ref1 = xp.where(cond, f(*arrays), fillvalue)
         ref2 = xp.where(cond, f(*arrays), f2(*arrays))
@@ -625,7 +625,7 @@ def f2(*args):
             # Python scalar. When it does, test can be run with array_api_strict, too.
             ref3 = xp.where(cond, f(*arrays), float_fillvalue)
 
-        if xp == np:  # because we ensured arrays are at least 1d
+        if is_numpy(xp):  # because we ensured arrays are at least 1d
             ref1 = ref1.reshape(result_shape)
             ref2 = ref2.reshape(result_shape)
             ref3 = ref3.reshape(result_shape)

scipy/_lib/tests/test_array_api.py

@@ -84,7 +84,7 @@ def test_strict_checks(self, xp, dtype, shape):
 
         kwarg_names = ["check_namespace", "check_dtype", "check_shape", "check_0d"]
         options = dict(zip(kwarg_names, [True, False, False, False]))
-        if xp == np:
+        if is_numpy(xp):
             xp_assert_equal(x, y, **options)
         else:
             with pytest.raises(

scipy/cluster/tests/test_vq.py

@@ -17,7 +17,7 @@
 
 from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api import (
-    SCIPY_ARRAY_API, array_namespace, xp_copy, xp_assert_close, xp_assert_equal
+    SCIPY_ARRAY_API, xp_copy, xp_assert_close, xp_assert_equal
 )
 
 skip_xp_backends = pytest.mark.skip_xp_backends
@@ -358,13 +358,12 @@ def test_krandinit(self, xp, krand_lock):
         datas = [xp.reshape(data, (200, 2)),
                  xp.reshape(data, (20, 20))[:10, :]]
         k = int(1e6)
-        xp_test = array_namespace(data)
         with krand_lock:
             for data in datas:
                 rng = np.random.default_rng(1234)
-                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)
+                init = _krandinit(data, k, rng, xp)
+                orig_cov = xpx.cov(data.T, xp=xp)
+                init_cov = xpx.cov(init.T, xp=xp)
                 xp_assert_close(orig_cov, init_cov, atol=1.1e-2)
 
     def test_kmeans2_empty(self, xp):

scipy/conftest.py

@@ -13,7 +13,7 @@
 
 from scipy._lib._fpumode import get_fpu_mode
 from scipy._lib._testutils import FPUModeChangeWarning
-from scipy._lib._array_api import SCIPY_ARRAY_API, SCIPY_DEVICE
+from scipy._lib._array_api import SCIPY_ARRAY_API, SCIPY_DEVICE, array_namespace
 from scipy._lib import _pep440
 
 try:
@@ -227,17 +227,21 @@ def xp(request):
     # if any, and raise pytest.xfail() if the current xp is in the list.
     skip_or_xfail_xp_backends(request, "xfail")
 
+    xp = request.param
+    # Potentially wrap namespace with array_api_compat
+    xp = array_namespace(xp.empty(0))
+
     if SCIPY_ARRAY_API:
         from scipy._lib._array_api import default_xp
 
         # Throughout all calls to assert_almost_equal, assert_array_almost_equal, and
         # xp_assert_* functions, test that the array namespace is xp in both the
         # expected and actual arrays. This is to detect the case where both arrays are
         # erroneously just plain numpy while xp is something else.
-        with default_xp(request.param):
-            yield request.param
+        with default_xp(xp):
+            yield xp
     else:
-        yield request.param
+        yield xp
 
 
 skip_xp_invalid_arg = pytest.mark.skipif(SCIPY_ARRAY_API,

scipy/differentiate/tests/test_differentiate.py

@@ -5,7 +5,7 @@
 
 import scipy._lib._elementwise_iterative_method as eim
 from scipy._lib._array_api_no_0d import xp_assert_close, xp_assert_equal, xp_assert_less
-from scipy._lib._array_api import is_numpy, is_torch, array_namespace
+from scipy._lib._array_api import is_numpy, is_torch, xp_size
 
 from scipy import stats, optimize, special
 from scipy.differentiate import derivative, jacobian, hessian
@@ -376,8 +376,7 @@ def test_special_cases(self, xp):
         # Test that integers are not passed to `f`
         # (otherwise this would overflow)
         def f(x):
-            xp_test = array_namespace(x)  # needs `isdtype`
-            assert xp_test.isdtype(x.dtype, 'real floating')
+            assert xp.isdtype(x.dtype, 'real floating')
             return x ** 99 - 1
 
         if not is_torch(xp):  # torch defaults to float32
@@ -582,9 +581,10 @@ def df1_1xy(x, y):
             return xp.sin(2*x) * y**2
 
         res = jacobian(df1, z, initial_step=10)
-        if is_numpy(xp):
-            assert len(np.unique(res.nit)) == 4
-            assert len(np.unique(res.nfev)) == 4
+        # FIXME https://github.com/scipy/scipy/pull/22320#discussion_r1914898175
+        if not is_torch(xp):
+            assert xp_size(xp.unique_values(res.nit)) == 4
+            assert xp_size(xp.unique_values(res.nfev)) == 4
 
         res00 = jacobian(lambda x: df1_0xy(x, z[1]), z[0:1], initial_step=10)
         res01 = jacobian(lambda y: df1_0xy(z[0], y), z[1:2], initial_step=10)
@@ -594,7 +594,10 @@ def df1_1xy(x, y):
         for attr in ['success', 'status', 'df', 'nit', 'nfev']:
             ref_attr = xp.asarray([[getattr(res00, attr), getattr(res01, attr)],
                                    [getattr(res10, attr), getattr(res11, attr)]])
-            ref[attr] = xp.squeeze(ref_attr)
+            ref[attr] = xp.squeeze(
+                ref_attr,
+                axis=tuple(ax for ax, size in enumerate(ref_attr.shape) if size == 1)
+            )            
             rtol = 1.5e-5 if res[attr].dtype == xp.float32 else 1.5e-14
             xp_assert_close(res[attr], ref[attr], rtol=rtol)
 
@@ -659,10 +662,9 @@ def test_float32(self, xp):
 
     def test_nfev(self, xp):
         z = xp.asarray([0.5, 0.25])
-        xp_test = array_namespace(z)
 
         def f1(z):
-            x, y = xp_test.broadcast_arrays(*z)
+            x, y = xp.broadcast_arrays(*z)
             f1.nfev = f1.nfev + (math.prod(x.shape[2:]) if x.ndim > 2 else 1)
             return xp.sin(x) * y ** 3
         f1.nfev = 0

scipy/fft/tests/test_basic.py

@@ -8,7 +8,7 @@
 from pytest import raises as assert_raises
 import scipy.fft as fft
 from scipy._lib._array_api import (
-    array_namespace, xp_size, xp_assert_close, xp_assert_equal
+    is_numpy, xp_size, xp_assert_close, xp_assert_equal
 )
 
 skip_xp_backends = pytest.mark.skip_xp_backends
@@ -228,10 +228,10 @@ def _check_axes(self, op, xp):
         dtype = get_expected_input_dtype(op, xp)
         x = xp.asarray(random((30, 20, 10)), dtype=dtype)
         axes = [(0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0)]
-        xp_test = array_namespace(x)
+
         for a in axes:
-            op_tr = op(xp_test.permute_dims(x, axes=a))
-            tr_op = xp_test.permute_dims(op(x, axes=a), axes=a)
+            op_tr = op(xp.permute_dims(x, axes=a))
+            tr_op = xp.permute_dims(op(x, axes=a), axes=a)
             xp_assert_close(op_tr, tr_op)
 
     @pytest.mark.parametrize("op", [fft.fftn, fft.ifftn, fft.rfftn, fft.irfftn])
@@ -248,15 +248,15 @@ def test_axes_subset_with_shape_standard(self, op, xp):
         dtype = get_expected_input_dtype(op, xp)
         x = xp.asarray(random((16, 8, 4)), dtype=dtype)
         axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
-        xp_test = array_namespace(x)
+
         for a in axes:
             # different shape on the first two axes
             shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
                            for ax in range(x.ndim)])
             # transform only the first two axes
-            op_tr = op(xp_test.permute_dims(x, axes=a),
+            op_tr = op(xp.permute_dims(x, axes=a),
                        s=shape[:2], axes=(0, 1))
-            tr_op = xp_test.permute_dims(op(x, s=shape[:2], axes=a[:2]),
+            tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]),
                                          axes=a)
             xp_assert_close(op_tr, tr_op)
 
@@ -268,21 +268,21 @@ def test_axes_subset_with_shape_non_standard(self, op, xp):
         dtype = get_expected_input_dtype(op, xp)
         x = xp.asarray(random((16, 8, 4)), dtype=dtype)
         axes = [(0, 1, 2), (0, 2, 1), (1, 2, 0)]
-        xp_test = array_namespace(x)
+
         for a in axes:
             # different shape on the first two axes
             shape = tuple([2*x.shape[ax] if ax in a[:2] else x.shape[ax]
                            for ax in range(x.ndim)])
             # transform only the first two axes
-            op_tr = op(xp_test.permute_dims(x, axes=a), s=shape[:2], axes=(0, 1))
-            tr_op = xp_test.permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a)
+            op_tr = op(xp.permute_dims(x, axes=a), s=shape[:2], axes=(0, 1))
+            tr_op = xp.permute_dims(op(x, s=shape[:2], axes=a[:2]), axes=a)
             xp_assert_close(op_tr, tr_op)
 
     def test_all_1d_norm_preserving(self, xp):
         # verify that round-trip transforms are norm-preserving
         x = xp.asarray(random(30), dtype=xp.float64)
-        xp_test = array_namespace(x)
-        x_norm = xp_test.linalg.vector_norm(x)
+
+        x_norm = xp.linalg.vector_norm(x)
         n = xp_size(x) * 2
         func_pairs = [(fft.rfft, fft.irfft),
                       # hfft: order so the first function takes x.size samples
@@ -294,12 +294,12 @@ def test_all_1d_norm_preserving(self, xp):
         for forw, back in func_pairs:
             if forw == fft.fft:
                 x = xp.asarray(x, dtype=xp.complex128)
-                x_norm = xp_test.linalg.vector_norm(x)
+                x_norm = xp.linalg.vector_norm(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)
-                    xp_assert_close(xp_test.linalg.vector_norm(tmp), x_norm)
+                    xp_assert_close(xp.linalg.vector_norm(tmp), x_norm)
 
     @skip_xp_backends(np_only=True)
     @pytest.mark.parametrize("dtype", [np.float16, np.longdouble])
@@ -481,13 +481,17 @@ def test_non_standard_params(func, xp):
     else:
         dtype = xp.complex128
 
-    if xp.__name__ != 'numpy':
-        x = xp.asarray([1, 2, 3], dtype=dtype)
-        # func(x) should not raise an exception
-        func(x)
+    x = xp.asarray([1, 2, 3], dtype=dtype)
+    # func(x) should not raise an exception
+    func(x)
+
+    if is_numpy(xp):
+        func(x, workers=2)
+    else:
         assert_raises(ValueError, func, x, workers=2)
-        # `plan` param is not tested since SciPy does not use it currently
-        # but should be tested if it comes into use
+
+    # `plan` param is not tested since SciPy does not use it currently
+    # but should be tested if it comes into use
 
 
 @pytest.mark.parametrize("dtype", ['float32', 'float64'])

scipy/fft/tests/test_fftlog.py

@@ -7,7 +7,7 @@
 from scipy.fft._fftlog import fht, ifht, fhtoffset
 from scipy.special import poch
 
-from scipy._lib._array_api import xp_assert_close, xp_assert_less, array_namespace
+from scipy._lib._array_api import xp_assert_close, xp_assert_less
 
 skip_xp_backends = pytest.mark.skip_xp_backends
 
@@ -186,13 +186,12 @@ def test_array_like(xp, op):
 @pytest.mark.parametrize('n', [128, 129])
 def test_gh_21661(xp, n):
     one = xp.asarray(1.0)
-    xp_test = array_namespace(one)
     mu = 0.0
     r = np.logspace(-7, 1, n)
     dln = math.log(r[1] / r[0])
     offset = fhtoffset(dln, initial=-6 * np.log(10), mu=mu)
     r = xp.asarray(r, dtype=one.dtype)
-    k = math.exp(offset) / xp_test.flip(r, axis=-1)
+    k = math.exp(offset) / xp.flip(r, axis=-1)
 
     def f(x, mu):
         return x**(mu + 1)*xp.exp(-x**2/2)

scipy/fft/tests/test_helper.py

@@ -11,7 +11,7 @@
 import numpy as np
 import sys
 from scipy._lib._array_api import (
-    xp_assert_close, get_xp_devices, xp_device, array_namespace
+    xp_assert_close, get_xp_devices, xp_device
 )
 from scipy import fft
 
@@ -532,11 +532,10 @@ def test_definition(self, xp):
         xp_assert_close(y, x2, check_dtype=False)
 
     def test_device(self, xp):
-        xp_test = array_namespace(xp.empty(0))
         devices = get_xp_devices(xp)
         for d in devices:
             y = fft.fftfreq(9, xp=xp, device=d)
-            x = xp_test.empty(0, device=d)
+            x = xp.empty(0, device=d)
             assert xp_device(y) == xp_device(x)
 
 
@@ -565,9 +564,8 @@ def test_definition(self, xp):
         xp_assert_close(y, x2, check_dtype=False)
 
     def test_device(self, xp):
-        xp_test = array_namespace(xp.empty(0))
         devices = get_xp_devices(xp)
         for d in devices:
             y = fft.rfftfreq(9, xp=xp, device=d)
-            x = xp_test.empty(0, device=d)
+            x = xp.empty(0, device=d)
             assert xp_device(y) == xp_device(x)

scipy/integrate/tests/test_cubature.py

@@ -115,13 +115,10 @@ def genz_malik_1980_f_2_exact(a, b, alphas, betas, xp):
     a = xp.reshape(a, (*([1]*(len(alphas.shape) - 1)), ndim))
     b = xp.reshape(b, (*([1]*(len(alphas.shape) - 1)), ndim))
 
-    # `xp` is the unwrapped namespace, so `.atan` won't work for `xp = np` and np<2.
-    xp_test = array_namespace(a)
-
     return (
         (-1)**ndim * 1/xp.prod(alphas, axis=-1)
         * xp.prod(
-            xp_test.atan((a - betas)/alphas) - xp_test.atan((b - betas)/alphas),
+            xp.atan((a - betas)/alphas) - xp.atan((b - betas)/alphas),
             axis=-1,
         )
     )
@@ -1344,19 +1341,18 @@ def test_infinite_limits_maintains_points(self, a, b, points, xp):
         transformation.
         """
 
-        xp_compat = array_namespace(xp.empty(0))
         points = [xp.asarray(p, dtype=xp.float64) for p in points]
 
         f_transformed = _InfiniteLimitsTransform(
             # Bind `points` and `xp` argument in f
-            lambda x: f_with_problematic_points(x, points, xp_compat),
-            xp.asarray(a, dtype=xp_compat.float64),
-            xp.asarray(b, dtype=xp_compat.float64),
-            xp=xp_compat,
+            lambda x: f_with_problematic_points(x, points, xp),
+            xp.asarray(a, dtype=xp.float64),
+            xp.asarray(b, dtype=xp.float64),
+            xp=xp,
         )
 
         for point in points:
-            transformed_point = f_transformed.inv(xp_compat.reshape(point, (1, -1)))
+            transformed_point = f_transformed.inv(xp.reshape(point, (1, -1)))
 
             with pytest.raises(Exception, match="called with a problematic point"):
                 f_transformed(transformed_point)