ENH: array types: add dask.array support

Author: lucascolley

.github/workflows/array_api.yml

@@ -43,8 +43,8 @@ jobs:
     name: Get commit message
     uses: ./.github/workflows/commit_message.yml
 
-  pytorch_cpu:
-    name: Linux PyTorch/JAX/xp-strict CPU
+  xp_cpu:
+    name: Linux PyTorch/JAX/Dask/xp-strict CPU
     needs: get_commit_message
     if: >
       needs.get_commit_message.outputs.message == 1
@@ -86,6 +86,10 @@ jobs:
       run: |
         python -m pip install "jax[cpu]"
 
+    - name: Install Dask
+      run: |
+        python -m pip install git+https://github.com/dask/dask.git
+
     - name:  Prepare compiler cache
       id:    prep-ccache
       shell: bash

dev.py

@@ -710,7 +710,8 @@ class Test(Task):
         multiple=True,
         help=(
             "Array API backend "
-            "('all', 'numpy', 'torch', 'cupy', 'array_api_strict', 'jax.numpy')."
+            "('all', 'numpy', 'torch', 'cupy', 'array_api_strict',"
+            " 'jax.numpy', 'dask.array')."
         )
     )
     # Argument can't have `help=`; used to consume all of `-- arg1 arg2 arg3`

scipy/_lib/_array_api.py

@@ -24,6 +24,7 @@
     is_cupy_namespace as is_cupy,
     is_torch_namespace as is_torch,
     is_jax_namespace as is_jax,
+    is_dask_namespace as is_dask,
     is_array_api_strict_namespace as is_array_api_strict
 )
 
@@ -246,6 +247,9 @@ def _strict_check(actual, desired, xp, *,
         assert actual.dtype == desired.dtype, _msg
 
     if check_shape:
+        if is_dask(xp):
+            actual.compute_chunk_sizes()
+            desired.compute_chunk_sizes()
         _msg = f"Shapes do not match.\nActual: {actual.shape}\nDesired: {desired.shape}"
         assert actual.shape == desired.shape, _msg
 

scipy/_lib/tests/test_array_api.py

@@ -4,7 +4,7 @@
 from scipy.conftest import array_api_compatible
 from scipy._lib._array_api import (
     _GLOBAL_CONFIG, array_namespace, _asarray, xp_copy, xp_assert_equal, is_numpy,
-    np_compat,
+    np_compat, is_dask
 )
 from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api_no_0d import xp_assert_equal as xp_assert_equal_no_0d
@@ -77,6 +77,10 @@ def test_copy(self, xp):
             x[1] = 11
             x[2] = 12
 
+            if is_dask(xp):
+                x.compute()
+                y.compute()
+
             assert x[0] != y[0]
             assert x[1] != y[1]
             assert x[2] != y[2]

scipy/conftest.py

@@ -12,7 +12,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, xp_device
 from scipy._lib import _pep440
 
 try:
@@ -178,6 +178,12 @@ def num_parallel_threads():
     except ImportError:
         pass
 
+    try:
+        import dask.array  # type: ignore[import-not-found]
+        xp_available_backends.update({'dask.array': dask.array})
+    except ImportError:
+        pass
+
     # by default, use all available backends
     if SCIPY_ARRAY_API.lower() not in ("1", "true"):
         SCIPY_ARRAY_API_ = json.loads(SCIPY_ARRAY_API)
@@ -366,6 +372,9 @@ def skip_or_xfail_xp_backends(xp, backends, kwargs, skip_or_xfail='skip'):
                 for d in xp.empty(0).devices():
                     if 'cpu' not in d.device_kind:
                         skip_or_xfail(reason=reason)
+            elif xp.__name__ == 'dask.array':
+                if xp_device(xp.empty(0)) != 'cpu':
+                    skip_or_xfail(reason=reason)
 
 
 # Following the approach of NumPy's conftest.py...

scipy/ndimage/tests/test_filters.py

@@ -195,6 +195,7 @@ def test_correlate01(self, xp):
 
     @xfail_xp_backends('cupy', reason="Differs by a factor of two?")
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_correlate01_overlap(self, xp):
         array = xp.reshape(xp.arange(256), (16, 16))
         weights = xp.asarray([2])
@@ -537,6 +538,7 @@ def test_correlate22(self, dtype_array, dtype_output, xp):
         assert_array_almost_equal(output, expected)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype_array', types)
     @pytest.mark.parametrize('dtype_output', types)
     def test_correlate23(self, dtype_array, dtype_output, xp):
@@ -556,6 +558,7 @@ def test_correlate23(self, dtype_array, dtype_output, xp):
         assert_array_almost_equal(output, expected)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype_array', types)
     @pytest.mark.parametrize('dtype_output', types)
     def test_correlate24(self, dtype_array, dtype_output, xp):
@@ -576,6 +579,7 @@ def test_correlate24(self, dtype_array, dtype_output, xp):
         assert_array_almost_equal(output, tcov)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype_array', types)
     @pytest.mark.parametrize('dtype_output', types)
     def test_correlate25(self, dtype_array, dtype_output, xp):
@@ -881,6 +885,7 @@ def test_gauss06(self, xp):
         assert_array_almost_equal(output1, output2)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_gauss_memory_overlap(self, xp):
         input = xp.arange(100 * 100, dtype=xp.float32)
         input = xp.reshape(input, (100, 100))
@@ -1227,6 +1232,7 @@ def test_prewitt01(self, dtype, xp):
         assert_array_almost_equal(t, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype', types + complex_types)
     def test_prewitt02(self, dtype, xp):
         if is_torch(xp) and dtype in ("uint16", "uint32", "uint64"):
@@ -1289,6 +1295,7 @@ def test_sobel01(self, dtype, xp):
         assert_array_almost_equal(t, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.",)
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype', types + complex_types)
     def test_sobel02(self, dtype, xp):
         if is_torch(xp) and dtype in ("uint16", "uint32", "uint64"):
@@ -1349,6 +1356,7 @@ def test_laplace01(self, dtype, xp):
         assert_array_almost_equal(tmp1 + tmp2, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only",)
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype',
                              ["int32", "float32", "float64",
                               "complex64", "complex128"])
@@ -1379,6 +1387,7 @@ def test_gaussian_laplace01(self, dtype, xp):
         assert_array_almost_equal(tmp1 + tmp2, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype',
                              ["int32", "float32", "float64",
                               "complex64", "complex128"])
@@ -1395,6 +1404,7 @@ def test_gaussian_laplace02(self, dtype, xp):
         assert_array_almost_equal(tmp1 + tmp2, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype', types + complex_types)
     def test_generic_laplace01(self, dtype, xp):
         if is_torch(xp) and dtype in ("uint16", "uint32", "uint64"):
@@ -1420,6 +1430,7 @@ def derivative2(input, axis, output, mode, cval, a, b):
         assert_array_almost_equal(tmp, output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype',
                              ["int32", "float32", "float64",
                               "complex64", "complex128"])
@@ -1441,6 +1452,7 @@ def test_gaussian_gradient_magnitude01(self, dtype, xp):
         xp_assert_close(output, expected, rtol=1e-6, atol=1e-6)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @pytest.mark.parametrize('dtype',
                              ["int32", "float32", "float64",
                               "complex64", "complex128"])
@@ -2640,6 +2652,7 @@ def test_gaussian_radius_invalid(xp):
 
 
 @skip_xp_backends("jax.numpy", reason="output array is read-only")
+@skip_xp_backends("dask.array", reason="output array is read-only.")
 class TestThreading:
     def check_func_thread(self, n, fun, args, out):
         from threading import Thread

scipy/ndimage/tests/test_morphology.py

@@ -2303,6 +2303,7 @@ def test_grey_erosion01(self, xp):
                                               [5, 5, 3, 3, 1]]))
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     @xfail_xp_backends("cupy", reason="https://github.com/cupy/cupy/issues/8398")
     def test_grey_erosion01_overlap(self, xp):
 
@@ -2498,6 +2499,7 @@ def test_morphological_laplace02(self, xp):
         assert_array_almost_equal(output, expected)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_white_tophat01(self, xp):
         array = xp.asarray([[3, 2, 5, 1, 4],
                             [7, 6, 9, 3, 5],
@@ -2551,6 +2553,7 @@ def test_white_tophat03(self, xp):
         xp_assert_equal(output, expected)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_white_tophat04(self, xp):
         array = np.eye(5, dtype=bool)
         structure = np.ones((3, 3), dtype=bool)
@@ -2563,6 +2566,7 @@ def test_white_tophat04(self, xp):
         ndimage.white_tophat(array, structure=structure, output=output)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_black_tophat01(self, xp):
         array = xp.asarray([[3, 2, 5, 1, 4],
                             [7, 6, 9, 3, 5],
@@ -2616,6 +2620,7 @@ def test_black_tophat03(self, xp):
         xp_assert_equal(output, expected)
 
     @skip_xp_backends("jax.numpy", reason="output array is read-only.")
+    @skip_xp_backends("dask.array", reason="output array is read-only.")
     def test_black_tophat04(self, xp):
         array = xp.asarray(np.eye(5, dtype=bool))
         structure = xp.asarray(np.ones((3, 3), dtype=bool))

scipy/special/tests/test_support_alternative_backends.py

@@ -5,7 +5,7 @@
 from scipy.conftest import array_api_compatible
 from scipy import special
 from scipy._lib._array_api_no_0d import xp_assert_close
-from scipy._lib._array_api import is_jax, is_torch, SCIPY_DEVICE
+from scipy._lib._array_api import is_jax, is_torch, SCIPY_DEVICE, is_dask
 from scipy._lib.array_api_compat import numpy as np
 
 try:
@@ -64,6 +64,9 @@ def test_support_alternative_backends(xp, f_name_n_args, dtype, shapes):
     ):
         pytest.skip(f"`{f_name}` does not have an array-agnostic implementation "
                     f"and cannot delegate to PyTorch.")
+        
+    if is_dask(xp) and f_name == 'rel_entr':
+        pytest.skip("boolean index assignment")
 
     shapes = shapes[:n_args]
     f = getattr(special, f_name)

scipy/stats/tests/test_entropy.py

@@ -11,8 +11,11 @@
 from scipy._lib._array_api_no_0d import (xp_assert_close, xp_assert_equal,
                                          xp_assert_less)
 
+@pytest.mark.skip_xp_backends("dask.array", reason="boolean index assignment")
+@pytest.mark.usefixtures("skip_xp_backends")
+@array_api_compatible
 class TestEntropy:
-    @array_api_compatible
+
     def test_entropy_positive(self, xp):
         # See ticket #497
         pk = xp.asarray([0.5, 0.2, 0.3])
@@ -22,7 +25,6 @@ def test_entropy_positive(self, xp):
         xp_assert_equal(eself, xp.asarray(0.))
         xp_assert_less(-edouble, xp.asarray(0.))
 
-    @array_api_compatible
     def test_entropy_base(self, xp):
         pk = xp.ones(16)
         S = stats.entropy(pk, base=2.)
@@ -34,21 +36,18 @@ def test_entropy_base(self, xp):
         S2 = stats.entropy(pk, qk, base=2.)
         xp_assert_less(xp.abs(S/S2 - math.log(2.)), xp.asarray(1.e-5))
 
-    @array_api_compatible
     def test_entropy_zero(self, xp):
         # Test for PR-479
         x = xp.asarray([0., 1., 2.])
         xp_assert_close(stats.entropy(x),
                         xp.asarray(0.63651416829481278))
 
-    @array_api_compatible
     def test_entropy_2d(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]])
         xp_assert_close(stats.entropy(pk, qk),
                         xp.asarray([0.1933259, 0.18609809]))
 
-    @array_api_compatible
     def test_entropy_2d_zero(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.0, 0.1], [0.3, 0.6], [0.5, 0.3]])
@@ -59,54 +58,46 @@ def test_entropy_2d_zero(self, xp):
         xp_assert_close(stats.entropy(pk, qk),
                         xp.asarray([0.17403988, 0.18609809]))
 
-    @array_api_compatible
     def test_entropy_base_2d_nondefault_axis(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         xp_assert_close(stats.entropy(pk, axis=1),
                         xp.asarray([0.63651417, 0.63651417, 0.66156324]))
 
-    @array_api_compatible
     def test_entropy_2d_nondefault_axis(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]])
         xp_assert_close(stats.entropy(pk, qk, axis=1),
                         xp.asarray([0.23104906, 0.23104906, 0.12770641]))
 
-    @array_api_compatible
     def test_entropy_raises_value_error(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.1, 0.2], [0.6, 0.3]])
         message = "Array shapes are incompatible for broadcasting."
         with pytest.raises(ValueError, match=message):
             stats.entropy(pk, qk)
 
-    @array_api_compatible
     def test_base_entropy_with_axis_0_is_equal_to_default(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         xp_assert_close(stats.entropy(pk, axis=0),
                         stats.entropy(pk))
 
-    @array_api_compatible
     def test_entropy_with_axis_0_is_equal_to_default(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]])
         xp_assert_close(stats.entropy(pk, qk, axis=0),
                         stats.entropy(pk, qk))
 
-    @array_api_compatible
     def test_base_entropy_transposed(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         xp_assert_close(stats.entropy(pk.T),
                         stats.entropy(pk, axis=1))
 
-    @array_api_compatible
     def test_entropy_transposed(self, xp):
         pk = xp.asarray([[0.1, 0.2], [0.6, 0.3], [0.3, 0.5]])
         qk = xp.asarray([[0.2, 0.1], [0.3, 0.6], [0.5, 0.3]])
         xp_assert_close(stats.entropy(pk.T, qk.T),
                         stats.entropy(pk, qk, axis=1))
 
-    @array_api_compatible
     def test_entropy_broadcasting(self, xp):
         rng = np.random.default_rng(74187315492831452)
         x = xp.asarray(rng.random(3))
@@ -115,22 +106,21 @@ def test_entropy_broadcasting(self, xp):
         xp_assert_equal(res[0], stats.entropy(x, y[0, ...]))
         xp_assert_equal(res[1], stats.entropy(x, y[1, ...]))
 
-    @array_api_compatible
     def test_entropy_shape_mismatch(self, xp):
         x = xp.ones((10, 1, 12))
         y = xp.ones((11, 2))
         message = "Array shapes are incompatible for broadcasting."
         with pytest.raises(ValueError, match=message):
             stats.entropy(x, y)
 
-    @array_api_compatible
     def test_input_validation(self, xp):
         x = xp.ones(10)
         message = "`base` must be a positive number."
         with pytest.raises(ValueError, match=message):
             stats.entropy(x, base=-2)
 
 
+@pytest.mark.skip_xp_backends("dask.array", reason="No sorting in Dask")
 @array_api_compatible
 @pytest.mark.usefixtures("skip_xp_backends")
 class TestDifferentialEntropy: