Fix dask

Author: lithomas1

.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

@@ -179,7 +179,9 @@ def device(x: Array, /) -> Device:
     out: device
         a ``device`` object (see the "Device Support" section of the array API specification).
     """
-    if is_numpy_array(x):
+    if is_numpy_array(x) or is_dask_array(x):
+        # TODO: dask technically can support GPU arrays
+        # Detecting the array backend isn't easy for dask, though, so just return CPU for now
         return "cpu"
     if is_jax_array(x):
         # JAX has .device() as a method, but it is being deprecated so that it

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,

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