Fix dask

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# macOS specific iles
+.DS_Store

README.md

@@ -125,11 +125,11 @@ part of the specification but which are useful for using the array API:
   [`x.device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.device.html)
   in the array API specification. Included because `numpy.ndarray` does not
   include the `device` attribute and this library does not wrap or extend the
-  array object. Note that for NumPy, `device(x)` is always `"cpu"`.
+  array object. Note that for NumPy and dask, `device(x)` is always `"cpu"`.
 
 - `to_device(x, device, /, *, stream=None)`: Equivalent to
   [`x.to_device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.to_device.html).
-  Included because neither NumPy's, CuPy's, nor PyTorch's array objects
+  Included because neither NumPy's, CuPy's, Dask's, nor PyTorch's array objects
   include this method. For NumPy, this function effectively does nothing since
   the only supported device is the CPU, but for CuPy, this method supports
   CuPy CUDA
@@ -240,6 +240,30 @@ Unlike the other libraries supported here, JAX array API support is contained
 entirely in the JAX library. The JAX array API support is tracked at
 https://github.com/google/jax/issues/18353.
 
+## Dask
+
+If you're using dask with numpy, many of the same limitations that apply to numpy
+will also apply to dask. Besides those differences, other limitations include missing
+sort functionality (no `sort` or `argsort`), and limited support for the optional `linalg`
+and `fft` extensions.
+
+In particular, the `fft` namespace is not compliant with the array API spec. Any functions
+that you find under the `fft` namespace are the original, unwrapped functions under [`dask.array.fft`](https://docs.dask.org/en/latest/array-api.html#fast-fourier-transforms), which may or may not be Array API compliant. Use at your own risk!
+
+For `linalg`, several methods are missing, for example:
+- `cross`
+- `det`
+- `eigh`
+- `eigvalsh`
+- `matrix_power`
+- `pinv`
+- `slogdet`
+- `matrix_norm`
+- `matrix_rank`
+Other methods may only be partially implemented or return incorrect results at times.
+
+The minimum supported Dask version is 2023.12.0.
+
 ## Vendoring
 
 This library supports vendoring as an installation method. To vendor the

array_api_compat/common/_helpers.py

@@ -159,7 +159,16 @@ def _check_device(xp, device):
         if device not in ["cpu", None]:
             raise ValueError(f"Unsupported device for NumPy: {device!r}")
 
-# device() is not on numpy.ndarray and to_device() is not on numpy.ndarray
+# Placeholder object to represent the dask device
+# when the array backend is not the CPU.
+# (since it is not easy to tell which device a dask array is on)
+class _dask_device:
+    def __repr__(self):
+        return "DASK_DEVICE"
+
+DASK_DEVICE = _dask_device()
+
+# device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray
 # or cupy.ndarray. They are not included in array objects of this library
 # because this library just reuses the respective ndarray classes without
 # wrapping or subclassing them. These helper functions can be used instead of
@@ -181,7 +190,17 @@ def device(x: Array, /) -> Device:
     """
     if is_numpy_array(x):
         return "cpu"
-    if is_jax_array(x):
+    elif is_dask_array(x):
+        # Peek at the metadata of the jax array to determine type
+        try:
+            import numpy as np
+            if isinstance(x._meta, np.ndarray):
+                # Must be on CPU since backed by numpy
+                return "cpu"
+        except ImportError:
+            pass
+        return DASK_DEVICE
+    elif is_jax_array(x):
         # JAX has .device() as a method, but it is being deprecated so that it
         # can become a property, in accordance with the standard. In order for
         # this function to not break when JAX makes the flip, we check for

array_api_compat/dask/array/__init__.py

@@ -22,7 +22,7 @@
 from dask.array import (
     arctanh as atanh,
 )
-from dask.array import (
+from numpy import (
     bool_ as bool,
 )
 from dask.array import (
@@ -67,15 +67,15 @@
     uint64,
 )
 
-from ..common._helpers import (
+from ...common._helpers import (
     array_namespace,
     device,
     get_namespace,
     is_array_api_obj,
     size,
     to_device,
 )
-from ..internal import _get_all_public_members
+from ..._internal import _get_all_public_members
 from ._aliases import (
     UniqueAllResult,
     UniqueCountsResult,

array_api_compat/dask/array/linalg.py

@@ -7,7 +7,7 @@
 )
 from dask.array.linalg import *  # noqa: F401, F403
 
-from .._internal import _get_all_public_members
+from ..._internal import _get_all_public_members
 from ._aliases import (
     EighResult,
     QRResult,

setup.py

@@ -1,4 +1,4 @@
-from setuptools import setup, find_packages
+from setuptools import setup, find_namespace_packages
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
@@ -8,7 +8,7 @@
 setup(
     name='array_api_compat',
     version=array_api_compat.__version__,
-    packages=find_packages(include=['array_api_compat*']),
+    packages=find_namespace_packages(include=["array_api_compat*"]),
     author="Consortium for Python Data API Standards",
     description="A wrapper around NumPy and other array libraries to make them compatible with the Array API standard",
     long_description=long_description,

tests/_helpers.py

@@ -8,7 +8,7 @@
 def import_(library, wrapper=False):
     if library == 'cupy':
         return pytest.importorskip(library)
-    if 'jax' in library and sys.version_info <= (3, 8):
+    if 'jax' in library and sys.version_info < (3, 9):
         pytest.skip('JAX array API support does not support Python 3.8')
 
     if wrapper:

tests/test_common.py

@@ -31,9 +31,6 @@ def test_is_xp_array(library, func):
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
 def test_device(library):
-    if library == "dask.array":
-        pytest.xfail("device() needs to be fixed for dask")
-
     xp = import_(library, wrapper=True)
 
     # We can't test much for device() and to_device() other than that