fix arange test

Author: lithomas1

array_api_compat/_dask_ci_shim.py

@@ -1,13 +1,15 @@
 """
 A little CI shim for the dask backend that
 disables the dask scheduler
+
+It also lets you see the dask dashboard for debugging
+at http://127.0.0.1:8787/status
 """
 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/_helpers.py

@@ -47,7 +47,6 @@ def _is_dask_array(x):
 
     import dask.array
 
-    # TODO: Should we reject ndarray subclasses?
     return isinstance(x, dask.array.Array)
 
 def is_array_api_obj(x):

array_api_compat/dask/_aliases.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 from ..common import _aliases
+from ..common._helpers import _check_device
 
 from .._internal import get_xp
 
@@ -30,13 +33,44 @@
     result_type,
 )
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from typing import Optional, Union
+    from ..common._typing import ndarray, Device, Dtype
+
 import dask.array as da
 
 isdtype = get_xp(np)(_aliases.isdtype)
 astype = _aliases.astype
 
 # Common aliases
-arange = get_xp(da)(_aliases.arange)
+
+# This arange func is modified from the common one to
+# not pass stop/step as keyword arguments, which will cause
+# an error with dask
+def dask_arange(
+    start: Union[int, float],
+    /,
+    stop: Optional[Union[int, float]] = None,
+    step: Union[int, float] = 1,
+    *,
+    xp,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    **kwargs
+) -> ndarray:
+    _check_device(xp, device)
+    args = [start]
+    if stop is not None:
+        args.append(stop)
+    else:
+        # stop is None, so start is actually stop
+        # prepend the default value for start which is 0
+        args.insert(0, 0)
+    args.append(step)
+    return xp.arange(*args, dtype=dtype, **kwargs)
+
+arange = get_xp(da)(dask_arange)
 eye = get_xp(da)(_aliases.eye)
 
 from functools import partial

dask-skips.txt

@@ -1,5 +1,2 @@
 # 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

dask-xfails.txt

@@ -25,7 +25,8 @@ array_api_tests/test_creation_functions.py::test_eye
 # finfo(float32).eps returns float32 but should return float
 array_api_tests/test_data_type_functions.py::test_finfo[float32]
 
-# shape mismatch
+# out[-1]=dask.aray<getitem ...> but should be some floating number
+# (I think the test is not forcing the op to be computed?)
 array_api_tests/test_creation_functions.py::test_linspace
 
 # out=-0, but should be +0

tests/test_vendoring.py

@@ -18,6 +18,6 @@ def test_vendoring_torch():
     from vendor_test import uses_torch
     uses_torch._test_torch()
 
-def test_vendoring_torch():
-    from vendor_test import uses_torch
-    uses_torch._test_torch()
+def test_vendoring_dask():
+    from vendor_test import uses_dask
+    uses_dask._test_dask()

vendor_test/uses_dask.py

@@ -5,15 +5,15 @@
 import dask.array as da
 import numpy as np
 
-def _test_numpy():
+def _test_dask():
     a = dask_compat.asarray([1., 2., 3.])
     b = dask_compat.arange(3, dtype=dask_compat.float32)
 
     # np.pow does not exist. Update this to use something else if it is added
     res = dask_compat.pow(a, b)
     assert res.dtype == dask_compat.float64 == np.float64
-    assert isinstance(a, da.array)
-    assert isinstance(b, da.array)
-    assert isinstance(res, da.array)
+    assert isinstance(a, da.Array)
+    assert isinstance(b, da.Array)
+    assert isinstance(res, da.Array)
 
     np.testing.assert_allclose(res, [1., 2., 9.])