Add dask to array-api-compat

.github/workflows/array-api-tests-dask.yml

@@ -0,0 +1,11 @@
+name: Array API Tests (Dask)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-dask:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: dask
+      extra-requires: numpy
+      pytest-extra-args: --disable-deadline

.github/workflows/array-api-tests.yml

@@ -6,6 +6,9 @@ on:
       package-name:
         required: true
         type: string
+      extra-requires:
+        required: false
+        type: string
       package-version:
         required: false
         type: string
@@ -54,7 +57,7 @@ jobs:
       if: "! (matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21'))"
       run: |
         python -m pip install --upgrade pip
-        python -m pip install '${{ inputs.package-name }} ${{ inputs.package-version }}'
+        python -m pip install '${{ inputs.package-name }} ${{ inputs.package-version }}' ${{ inputs.extra-requires }}
         python -m pip install -r ${GITHUB_WORKSPACE}/array-api-tests/requirements.txt
     - name: Run the array API testsuite (${{ inputs.package-name }})
       if: "! (matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21'))"

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Install Dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest numpy torch
+          python -m pip install pytest numpy torch dask[array]
 
       - name: Run Tests
         run: |

array_api_compat/_dask_ci_shim.py

@@ -0,0 +1,13 @@
+"""
+A little CI shim for the dask backend that
+disables the dask scheduler
+"""
+import dask
+dask.config.set(scheduler='synchronous')
+
+from dask.distributed import Client
+_client = Client()
+print(_client.dashboard_link)
+
+from .dask import *
+from .dask import __array_api_version__

array_api_compat/common/_aliases.py

@@ -303,6 +303,8 @@ def _asarray(
         import numpy as xp
     elif namespace == 'cupy':
         import cupy as xp
+    elif namespace == 'dask':
+        import dask.array as xp
     else:
         raise ValueError("Unrecognized namespace argument to asarray()")
 
@@ -322,10 +324,13 @@ def _asarray(
     if copy in COPY_FALSE:
         # copy=False is not yet implemented in xp.asarray
         raise NotImplementedError("copy=False is not yet implemented")
-    if isinstance(obj, xp.ndarray):
+    # TODO: This feels wrong (__array__ is not in the standard)
+    # Dask doesn't support DLPack, though, so, this'll do
+    if (hasattr(xp, "ndarray") and isinstance(obj, xp.ndarray)) or hasattr(obj, "__array__"):
         if dtype is not None and obj.dtype != dtype:
             copy = True
-        if copy in COPY_TRUE:
+        # Dask arrays are immutable, so copy doesn't do anything
+        if copy in COPY_TRUE and namespace != "dask":
             return xp.array(obj, copy=True, dtype=dtype)
         return obj
 

array_api_compat/common/_helpers.py

@@ -40,6 +40,16 @@ def _is_torch_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, torch.Tensor)
 
+def _is_dask_array(x):
+    # Avoid importing dask if it isn't already
+    if 'dask.array' not in sys.modules:
+        return False
+
+    import dask.array
+
+    # TODO: Should we reject ndarray subclasses?
+    return isinstance(x, dask.array.Array)
+
 def is_array_api_obj(x):
     """
     Check if x is an array API compatible array object.
@@ -97,6 +107,13 @@ def your_function(x, y):
             else:
                 import torch
                 namespaces.add(torch)
+        elif _is_dask_array(x):
+            _check_api_version(api_version)
+            if _use_compat:
+                from .. import dask as dask_namespace
+                namespaces.add(dask_namespace)
+            else:
+                raise TypeError("_use_compat cannot be False if input array is a dask array!")
         else:
             # TODO: Support Python scalars?
             raise TypeError("The input is not a supported array type")
@@ -219,6 +236,13 @@ def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, A
         return _cupy_to_device(x, device, stream=stream)
     elif _is_torch_array(x):
         return _torch_to_device(x, device, stream=stream)
+    elif _is_dask_array(x):
+        if stream is not None:
+            raise ValueError("The stream argument to to_device() is not supported")
+        # TODO: What if our array is on the GPU already?
+        if device == 'cpu':
+            return x
+        raise ValueError(f"Unsupported device {device!r}")
     return x.to_device(device, stream=stream)
 
 def size(x):

array_api_compat/dask/__init__.py

@@ -0,0 +1,6 @@
+from dask.array import *
+
+# These imports may overwrite names from the import * above.
+from ._aliases import *
+
+__array_api_version__ = '2022.12'

array_api_compat/dask/_aliases.py

@@ -0,0 +1,86 @@
+from ..common import _aliases
+
+from .._internal import get_xp
+
+import numpy as np
+from numpy import (
+    # Constants
+    e,
+    inf,
+    nan,
+    pi,
+    newaxis,
+    # Dtypes
+    bool_ as bool,
+    float32,
+    float64,
+    int8,
+    int16,
+    int32,
+    int64,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    complex64,
+    complex128,
+    iinfo,
+    finfo,
+    can_cast,
+    result_type,
+)
+
+import dask.array as da
+
+isdtype = get_xp(np)(_aliases.isdtype)
+astype = _aliases.astype
+
+# Common aliases
+arange = get_xp(da)(_aliases.arange)
+eye = get_xp(da)(_aliases.eye)
+
+from functools import partial
+asarray = partial(_aliases._asarray, namespace='dask')
+asarray.__doc__ = _aliases._asarray.__doc__
+
+linspace = get_xp(da)(_aliases.linspace)
+eye = get_xp(da)(_aliases.eye)
+UniqueAllResult = get_xp(da)(_aliases.UniqueAllResult)
+UniqueCountsResult = get_xp(da)(_aliases.UniqueCountsResult)
+UniqueInverseResult = get_xp(da)(_aliases.UniqueInverseResult)
+unique_all = get_xp(da)(_aliases.unique_all)
+unique_counts = get_xp(da)(_aliases.unique_counts)
+unique_inverse = get_xp(da)(_aliases.unique_inverse)
+unique_values = get_xp(da)(_aliases.unique_values)
+permute_dims = get_xp(da)(_aliases.permute_dims)
+std = get_xp(da)(_aliases.std)
+var = get_xp(da)(_aliases.var)
+empty = get_xp(da)(_aliases.empty)
+empty_like = get_xp(da)(_aliases.empty_like)
+full = get_xp(da)(_aliases.full)
+full_like = get_xp(da)(_aliases.full_like)
+ones = get_xp(da)(_aliases.ones)
+ones_like = get_xp(da)(_aliases.ones_like)
+zeros = get_xp(da)(_aliases.zeros)
+zeros_like = get_xp(da)(_aliases.zeros_like)
+reshape = get_xp(da)(_aliases.reshape)
+matrix_transpose = get_xp(da)(_aliases.matrix_transpose)
+vecdot = get_xp(da)(_aliases.vecdot)
+
+from dask.array import (
+    # Element wise aliases
+    arccos as acos,
+    arccosh as acosh,
+    arcsin as asin,
+    arcsinh as asinh,
+    arctan as atan,
+    arctan2 as atan2,
+    arctanh as atanh,
+    left_shift as bitwise_left_shift,
+    right_shift as bitwise_right_shift,
+    invert as bitwise_invert,
+    power as pow,
+    # Other
+    concatenate as concat,
+)
+

array_api_compat/dask/linalg.py

@@ -0,0 +1,34 @@
+from dask.array.linalg import *
+from dask.array.linalg import __all__ as linalg_all
+
+from ..common import _linalg
+from .._internal import get_xp
+from ._aliases import (matmul, matrix_transpose, tensordot, vecdot)
+
+import dask.array as da
+
+cross = get_xp(da)(_linalg.cross)
+outer = get_xp(da)(_linalg.outer)
+EighResult = _linalg.EighResult
+QRResult = _linalg.QRResult
+SlogdetResult = _linalg.SlogdetResult
+SVDResult = _linalg.SVDResult
+eigh = get_xp(da)(_linalg.eigh)
+qr = get_xp(da)(_linalg.qr)
+slogdet = get_xp(da)(_linalg.slogdet)
+svd = get_xp(da)(_linalg.svd)
+cholesky = get_xp(da)(_linalg.cholesky)
+matrix_rank = get_xp(da)(_linalg.matrix_rank)
+pinv = get_xp(da)(_linalg.pinv)
+matrix_norm = get_xp(da)(_linalg.matrix_norm)
+svdvals = get_xp(da)(_linalg.svdvals)
+vector_norm = get_xp(da)(_linalg.vector_norm)
+diagonal = get_xp(da)(_linalg.diagonal)
+trace = get_xp(da)(_linalg.trace)
+
+__all__ = linalg_all + _linalg.__all__
+
+del get_xp
+del da
+del linalg_all
+del _linalg

dask-skips.txt

@@ -0,0 +1,5 @@
+# FFT isn't conformant
+array_api_tests/test_fft.py
+
+# Errors with dask, also makes Dask go OOM
+array_api_tests/test_creation_functions.py::test_arange