fft/linalg/special passing tests

Author: lithomas1

scipy/conftest.py

@@ -179,8 +179,11 @@ def num_parallel_threads():
         pass
 
     try:
-        import dask.array  # type: ignore[import-not-found]
-        xp_available_backends.update({'dask.array': dask.array})
+        # Note: dask.array main namespace is not array API compatible
+        # (to address this, we will fix tests that use the broken dask behavior to
+        # use the array-api-compat wrapped version instead)
+        import dask.array as da
+        xp_available_backends.update({'dask.array': da})
     except ImportError:
         pass
 

scipy/fft/tests/test_basic.py

@@ -333,8 +333,12 @@ def test_dtypes_real(self, dtype, xp):
 
     @pytest.mark.parametrize("dtype", ["complex64", "complex128"])
     def test_dtypes_complex(self, dtype, xp):
+        # Trick to get the array-api-compat namespace for dask
+        # (otherwise the "naked" dask.array asarray does not respect
+        # the input dtype)
+        xp_test = array_namespace(xp.asarray(1))
         rng = np.random.default_rng(1234)
-        x = xp.asarray(rng.random(30), dtype=getattr(xp, dtype))
+        x = xp.asarray(rng.random(30), dtype=getattr(xp_test, dtype))
 
         res_fft = fft.ifft(fft.fft(x))
         # Check both numerical results and exact dtype matches

scipy/fft/tests/test_real_transforms.py

@@ -7,7 +7,7 @@
 import scipy.fft as fft
 from scipy import fftpack
 from scipy.conftest import array_api_compatible
-from scipy._lib._array_api import xp_copy, xp_assert_close
+from scipy._lib._array_api import xp_copy, xp_assert_close, array_namespace
 
 pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends")]
 skip_xp_backends = pytest.mark.skip_xp_backends
@@ -199,6 +199,10 @@ def test_orthogonalize_noop(func, type, norm, xp):
 def test_orthogonalize_dct1(norm, xp):
     x = xp.asarray(np.random.rand(100))
 
+    # use array-api-compat namespace for dask
+    # since dask asarray never makes a copy
+    # which makes xp_copy silently a no-op
+    xp = array_namespace(x)
     x2 = xp_copy(x, xp=xp)
     x2[0] *= SQRT_2
     x2[-1] *= SQRT_2
@@ -232,6 +236,10 @@ def test_orthogonalize_dcst2(func, norm, xp):
 @pytest.mark.parametrize("func", [dct, dst])
 def test_orthogonalize_dcst3(func, norm, xp):
     x = xp.asarray(np.random.rand(100))
+    # use array-api-compat namespace for dask
+    # since dask asarray never makes a copy
+    # which makes xp_copy silently a no-op
+    xp = array_namespace(x)
     x2 = xp_copy(x, xp=xp)
     x2[0 if func == dct else -1] *= SQRT_2
 

scipy/special/_logsumexp.py

@@ -147,7 +147,9 @@ def _wrap_radians(x, xp=None):
     out = -((-x + math.pi) % (2 * math.pi) - math.pi)
     # preserve relative precision
     no_wrap = xp.abs(x) < xp.pi
-    out[no_wrap] = x[no_wrap]
+    # TODO: i think this is correct but double check
+    # out[no_wrap] = x[no_wrap]
+    out = xp.where(no_wrap, x, out)
     return out
 
 

scipy/special/tests/test_logsumexp.py

@@ -18,9 +18,6 @@
 
 
 @array_api_compatible
-@pytest.mark.usefixtures("skip_xp_backends")
-@pytest.mark.skip_xp_backends('jax.numpy',
-                              reason="JAX arrays do not support item assignment")
 def test_wrap_radians(xp):
     x = xp.asarray([-math.pi-1, -math.pi, -1, -1e-300,
                     0, 1e-300, 1, math.pi, math.pi+1])
@@ -35,6 +32,10 @@ def test_wrap_radians(xp):
 @pytest.mark.skip_xp_backends('jax.numpy',
                               reason="JAX arrays do not support item assignment")
 class TestLogSumExp:
+    # numpy warning filters don't work for dask
+    # (also we should not expect the numpy warning filter to work for any Array API
+    # library)
+    @pytest.mark.filterwarnings("ignore:divide by zero encountered in log")
     def test_logsumexp(self, xp):
         # Test with zero-size array
         a = xp.asarray([])
@@ -69,11 +70,8 @@ def test_logsumexp(self, xp):
         nan = xp.asarray([xp.nan])
         xp_assert_equal(logsumexp(inf), inf[0])
         xp_assert_equal(logsumexp(-inf), -inf[0])
-        # catch warnings here for dasks state there's no way to suppress
-        # warnings just for dask
-        # https://github.com/dask/dask/issues/3245
-        with np.errstate(divide='ignore', invalid='ignore'):
-            xp_assert_equal(logsumexp(nan), nan[0])
+
+        xp_assert_equal(logsumexp(nan), nan[0])
         xp_assert_equal(logsumexp(xp.asarray([-xp.inf, -xp.inf])), -inf[0])
 
         # Handling an array with different magnitudes on the axes
@@ -119,6 +117,7 @@ def test_logsumexp_sign(self, xp):
         xp_assert_close(r, xp.asarray(1.))
         xp_assert_equal(s, xp.asarray(-1.))
 
+    @pytest.mark.filterwarnings("ignore::RuntimeWarning")
     def test_logsumexp_sign_zero(self, xp):
         a = xp.asarray([1, 1])
         b = xp.asarray([1, -1])
@@ -223,6 +222,7 @@ def test_gh18295(self, xp):
         ref = xp.logaddexp(a[0], a[1])
         xp_assert_close(res, ref)
 
+    @pytest.mark.filterwarnings("ignore::FutureWarning:dask")
     @pytest.mark.parametrize('dtype', ['complex64', 'complex128'])
     def test_gh21610(self, xp, dtype):
         # gh-21610 noted that `logsumexp` could return imaginary components

scipy/special/tests/test_support_alternative_backends.py

@@ -52,6 +52,8 @@ def test_rel_entr_generic(dtype):
 # @pytest.mark.usefixtures("skip_xp_backends")
 # `reversed` is for developer convenience: test new function first = less waiting
 @pytest.mark.parametrize('f_name_n_args', reversed(array_special_func_map.items()))
+# numpy warning filter doesn't work for dask
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
 @pytest.mark.parametrize('dtype', ['float32', 'float64'])
 @pytest.mark.parametrize('shapes', [[(0,)]*4, [tuple()]*4, [(10,)]*4,
                                     [(10,), (11, 1), (12, 1, 1), (13, 1, 1, 1)]])