Merge pull request data-apis#120 from crusaderky/dask_nan_shape

MAINT: Upgrade to array-api-compat 1.11

Author: lucascolley

pixi.lock

@@ -26,7 +26,7 @@ classifiers = [
     "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["array-api-compat>=1.10.0,<2"]
+dependencies = ["array-api-compat>=1.11,<2"]
 
 [project.urls]
 Homepage = "https://github.com/data-apis/array-api-extra"
@@ -48,7 +48,7 @@ platforms = ["linux-64", "osx-arm64", "win-64"]
 
 [tool.pixi.dependencies]
 python = ">=3.10,<3.14"
-array-api-compat = ">=1.10.0,<2"
+array-api-compat = ">=1.11,<2"
 
 [tool.pixi.pypi-dependencies]
 array-api-extra = { path = ".", editable = true }

pyproject.toml

@@ -556,7 +556,7 @@ def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
     _, counts = xp.unique_counts(x)
     n = _compat.size(counts)
     # FIXME https://github.com/data-apis/array-api-compat/pull/231
-    if n is None or math.isnan(n):  # e.g. Dask, ndonnx
+    if n is None:  # e.g. Dask, ndonnx
         return xp.astype(counts, xp.bool).sum()
     return xp.asarray(n, device=_compat.device(x))
 

src/array_api_extra/_lib/_funcs.py

@@ -14,6 +14,7 @@
         is_dask_namespace,
         is_jax_array,
         is_jax_namespace,
+        is_lazy_array,
         is_numpy_array,
         is_numpy_namespace,
         is_pydata_sparse_array,
@@ -35,6 +36,7 @@
         is_dask_namespace,
         is_jax_array,
         is_jax_namespace,
+        is_lazy_array,
         is_numpy_array,
         is_numpy_namespace,
         is_pydata_sparse_array,
@@ -56,6 +58,7 @@
     "is_dask_namespace",
     "is_jax_array",
     "is_jax_namespace",
+    "is_lazy_array",
     "is_numpy_array",
     "is_numpy_namespace",
     "is_pydata_sparse_array",

src/array_api_extra/_lib/_utils/_compat.py

@@ -32,5 +32,6 @@ def is_jax_array(x: object, /) -> bool: ...
 def is_numpy_array(x: object, /) -> bool: ...
 def is_pydata_sparse_array(x: object, /) -> bool: ...
 def is_torch_array(x: object, /) -> bool: ...
+def is_lazy_array(x: object, /) -> bool: ...
 def is_writeable_array(x: object, /) -> bool: ...
 def size(x: Array, /) -> int | None: ...

src/array_api_extra/_lib/_utils/_compat.pyi

@@ -11,7 +11,7 @@
 from array_api_extra import at
 from array_api_extra._lib import Backend
 from array_api_extra._lib._at import _AtOp
-from array_api_extra._lib._testing import xfail, xp_assert_equal
+from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import array_namespace, is_writeable_array
 from array_api_extra._lib._utils._typing import Array, Index
 from array_api_extra.testing import lazy_xp_function
@@ -219,7 +219,6 @@ def test_alternate_index_syntax():
 def test_incompatible_dtype(
     xp: ModuleType,
     library: Backend,
-    request: pytest.FixtureRequest,
     op: _AtOp,
     copy: bool | None,
     bool_mask: bool,
@@ -253,8 +252,6 @@ def test_incompatible_dtype(
                 z = at_op(x, idx, op, 1.1, copy=copy)
 
     elif library is Backend.DASK:
-        if op in (_AtOp.MIN, _AtOp.MAX) and bool_mask:
-            xfail(request, reason="need array-api-compat 1.11")
         z = at_op(x, idx, op, 1.1, copy=copy)
 
     elif library is Backend.ARRAY_API_STRICT and op is not _AtOp.SET:

tests/test_at.py

@@ -32,16 +32,14 @@
 
 lazy_xp_function(atleast_nd, static_argnames=("ndim", "xp"))
 lazy_xp_function(cov, static_argnames="xp")
-# FIXME .device attribute https://github.com/data-apis/array-api-compat/pull/238
-lazy_xp_function(create_diagonal, jax_jit=False, static_argnames=("offset", "xp"))
+lazy_xp_function(create_diagonal, static_argnames=("offset", "xp"))
 lazy_xp_function(expand_dims, static_argnames=("axis", "xp"))
 lazy_xp_function(kron, static_argnames="xp")
 lazy_xp_function(nunique, static_argnames="xp")
 lazy_xp_function(pad, static_argnames=("pad_width", "mode", "constant_values", "xp"))
 # FIXME calls in1d which calls xp.unique_values without size
 lazy_xp_function(setdiff1d, jax_jit=False, static_argnames=("assume_unique", "xp"))
-# FIXME .device attribute https://github.com/data-apis/array-api-compat/pull/238
-lazy_xp_function(sinc, jax_jit=False, static_argnames="xp")
+lazy_xp_function(sinc, static_argnames="xp")
 
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")

tests/test_funcs.py

@@ -14,6 +14,7 @@ def test_vendor_compat():
         is_dask_namespace,
         is_jax_array,
         is_jax_namespace,
+        is_lazy_array,
         is_numpy_array,
         is_numpy_namespace,
         is_pydata_sparse_array,
@@ -35,6 +36,7 @@ def test_vendor_compat():
     assert not is_dask_namespace(xp)
     assert not is_jax_array(x)
     assert not is_jax_namespace(xp)
+    assert not is_lazy_array(x)
     assert not is_numpy_array(x)
     assert not is_numpy_namespace(xp)
     assert not is_pydata_sparse_array(x)