Merge branch 'main' into xp_assert_enhancements

Author: crusaderky

.github/workflows/docs-deploy.yml

@@ -1,7 +1,7 @@
 name: Docs Deploy
 
 permissions:
-  contents: read
+  contents: write # needed for the deploy step
 
 on:
   workflow_run:

codecov.yml

@@ -4,3 +4,6 @@ github_checks:
 ignore:
   - "src/array_api_extra/_lib/_compat"
   - "src/array_api_extra/_lib/_typing"
+coverage:
+  status:
+    project: off

pixi.lock

@@ -54,17 +54,17 @@ array-api-compat = ">=1.11.2,<2"
 array-api-extra = { path = ".", editable = true }
 
 [tool.pixi.feature.lint.dependencies]
-typing-extensions = ">=4.13.1"
+typing-extensions = ">=4.13.2"
 pre-commit = ">=4.2.0"
 pylint = ">=3.3.6"
 basedmypy = ">=2.10.0"
-basedpyright = ">=1.28.3"
+basedpyright = ">=1.28.5"
 numpydoc = ">=1.8.0,<2"
 # import dependencies for mypy:
 array-api-strict = ">=2.3.1"
 numpy = ">=2.1.3"
 pytest = ">=8.3.5"
-hypothesis = ">=6.130.11"
+hypothesis = ">=6.131.8"
 dask-core = ">=2025.3.0"     # No distributed, tornado, etc.
 # NOTE: don't add jax, pytorch, sparse, cupy here
 # as they slow down mypy and are not portable across target OSs
@@ -80,7 +80,7 @@ lint = { depends-on = ["pre-commit", "pylint", "mypy", "pyright"] , description
 [tool.pixi.feature.tests.dependencies]
 pytest = ">=8.3.5"
 pytest-cov = ">=6.1.1"
-hypothesis = ">=6.130.11"
+hypothesis = ">=6.131.8"
 array-api-strict = ">=2.3.1"
 numpy = ">=1.22.0"
 
@@ -107,7 +107,7 @@ sphinx-autodoc-typehints = ">=1.25.3"
 # Needed to import parsed modules with autodoc
 dask-core = ">=2025.3.0"
 pytest = ">=8.3.5"
-typing-extensions = ">=4.13.1"
+typing-extensions = ">=4.13.2"
 numpy = ">=2.1.3"
 
 [tool.pixi.feature.docs.tasks]
@@ -136,7 +136,7 @@ numpy = "=1.22.0"
 [tool.pixi.feature.backends.dependencies]
 pytorch = ">=2.6.0"
 dask = ">=2025.3.0"
-numba = ">=0.61.0"    # sparse dependency
+numba = ">=0.61.2"    # sparse dependency
 llvmlite = ">=0.44.0" # sparse dependency
 
 [tool.pixi.feature.backends.pypi-dependencies]

pyproject.toml

@@ -50,6 +50,22 @@
       "matchManagers": ["github-actions"],
       "matchPackageNames": ["python"],
       "enabled": false
+    },
+    {
+      "description": "Group Dask packages.",
+      "matchPackageNames": ["dask", "dask-core"],
+      "groupName": "dask"
+    },
+    {
+      "description": "Group JAX packages.",
+      "matchPackageNames": ["jax", "jaxlib"],
+      "groupName": "jax"
+    },
+    {
+      "description": "Schedule hypothesis monthly as releases are frequent.",
+      "matchManagers": ["pixi"],
+      "matchPackageNames": ["hypothesis"],
+      "schedule": ["* * 10 * *"]
     }
   ]
 }

renovate.json

@@ -23,6 +23,7 @@
         is_torch_namespace,
         is_writeable_array,
         size,
+        to_device,
     )
 except ImportError:
     from array_api_compat import (
@@ -45,6 +46,7 @@
         is_torch_namespace,
         is_writeable_array,
         size,
+        to_device,
     )
 
 __all__ = [
@@ -67,4 +69,5 @@
     "is_torch_namespace",
     "is_writeable_array",
     "size",
+    "to_device",
 ]

src/array_api_extra/_lib/_utils/_compat.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from types import ModuleType
+from typing import Any, TypeGuard
 
 # TODO import from typing (requires Python >=3.13)
 from typing_extensions import TypeIs
@@ -12,29 +13,33 @@ from ._typing import Array, Device
 
 # pylint: disable=missing-class-docstring,unused-argument
 
-class Namespace(ModuleType):
-    def device(self, x: Array, /) -> Device: ...
-
 def array_namespace(
     *xs: Array | complex | None,
     api_version: str | None = None,
     use_compat: bool | None = None,
-) -> Namespace: ...
+) -> ModuleType: ...
 def device(x: Array, /) -> Device: ...
 def is_array_api_obj(x: object, /) -> TypeIs[Array]: ...
-def is_array_api_strict_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_cupy_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_dask_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_jax_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_numpy_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_pydata_sparse_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_torch_namespace(xp: ModuleType, /) -> TypeIs[Namespace]: ...
-def is_cupy_array(x: object, /) -> TypeIs[Array]: ...
-def is_dask_array(x: object, /) -> TypeIs[Array]: ...
-def is_jax_array(x: object, /) -> TypeIs[Array]: ...
-def is_numpy_array(x: object, /) -> TypeIs[Array]: ...
-def is_pydata_sparse_array(x: object, /) -> TypeIs[Array]: ...
-def is_torch_array(x: object, /) -> TypeIs[Array]: ...
-def is_lazy_array(x: object, /) -> TypeIs[Array]: ...
-def is_writeable_array(x: object, /) -> TypeIs[Array]: ...
+def is_array_api_strict_namespace(xp: ModuleType, /) -> bool: ...
+def is_cupy_namespace(xp: ModuleType, /) -> bool: ...
+def is_dask_namespace(xp: ModuleType, /) -> bool: ...
+def is_jax_namespace(xp: ModuleType, /) -> bool: ...
+def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
+def is_pydata_sparse_namespace(xp: ModuleType, /) -> bool: ...
+def is_torch_namespace(xp: ModuleType, /) -> bool: ...
+def is_cupy_array(x: object, /) -> TypeGuard[Array]: ...
+def is_dask_array(x: object, /) -> TypeGuard[Array]: ...
+def is_jax_array(x: object, /) -> TypeGuard[Array]: ...
+def is_numpy_array(x: object, /) -> TypeGuard[Array]: ...
+def is_pydata_sparse_array(x: object, /) -> TypeGuard[Array]: ...
+def is_torch_array(x: object, /) -> TypeGuard[Array]: ...
+def is_lazy_array(x: object, /) -> TypeGuard[Array]: ...
+def is_writeable_array(x: object, /) -> TypeGuard[Array]: ...
 def size(x: Array, /) -> int | None: ...
+def to_device(  # type: ignore[explicit-any]
+    x: Array,
+    device: Device,  # pylint: disable=redefined-outer-name
+    /,
+    *,
+    stream: int | Any | None = None,
+) -> Array: ...

src/array_api_extra/_lib/_utils/_compat.pyi

@@ -39,7 +39,7 @@ def override(func: object) -> object:
 def lazy_xp_function(  # type: ignore[explicit-any]
     func: Callable[..., Any],
     *,
-    allow_dask_compute: int = 0,
+    allow_dask_compute: bool | int = False,
     jax_jit: bool = True,
     static_argnums: int | Sequence[int] | None = None,
     static_argnames: str | Iterable[str] | None = None,
@@ -59,9 +59,10 @@ def lazy_xp_function(  # type: ignore[explicit-any]
     ----------
     func : callable
         Function to be tested.
-    allow_dask_compute : int, optional
-        Number of times `func` is allowed to internally materialize the Dask graph. This
-        is typically triggered by ``bool()``, ``float()``, or ``np.asarray()``.
+    allow_dask_compute : bool | int, optional
+        Whether `func` is allowed to internally materialize the Dask graph, or maximum
+        number of times it is allowed to do so. This is typically triggered by
+        ``bool()``, ``float()``, or ``np.asarray()``.
 
         Set to 1 if you are aware that `func` converts the input parameters to NumPy and
         want to let it do so at least for the time being, knowing that it is going to be
@@ -75,7 +76,10 @@ def lazy_xp_function(  # type: ignore[explicit-any]
         a test function that invokes `func` multiple times should still work with this
         parameter set to 1.
 
-        Default: 0, meaning that `func` must be fully lazy and never materialize the
+        Set to True to allow `func` to materialize the graph an unlimited number
+        of times.
+
+        Default: False, meaning that `func` must be fully lazy and never materialize the
         graph.
     jax_jit : bool, optional
         Set to True to replace `func` with ``jax.jit(func)`` after calling the
@@ -235,6 +239,10 @@ def iter_tagged() -> (  # type: ignore[explicit-any]
     if is_dask_namespace(xp):
         for mod, name, func, tags in iter_tagged():
             n = tags["allow_dask_compute"]
+            if n is True:
+                n = 1_000_000
+            elif n is False:
+                n = 0
             wrapped = _dask_wrap(func, n)
             monkeypatch.setattr(mod, name, wrapped)
 

src/array_api_extra/testing.py

@@ -521,7 +521,7 @@ def test_xp(self, xp: ModuleType):
 class TestExpandDims:
     def test_single_axis(self, xp: ModuleType):
         """Trivial case where xpx.expand_dims doesn't add anything to xp.expand_dims"""
-        a = xp.empty((2, 3, 4, 5))
+        a = xp.asarray(np.reshape(np.arange(2 * 3 * 4 * 5), (2, 3, 4, 5)))
         for axis in range(-5, 4):
             b = expand_dims(a, axis=axis)
             xp_assert_equal(b, xp.expand_dims(a, axis=axis))

tests/test_funcs.py

@@ -200,13 +200,18 @@ def non_materializable4(x: Array) -> Array:
     return non_materializable(x)
 
 
+def non_materializable5(x: Array) -> Array:
+    return non_materializable(x)
+
+
 lazy_xp_function(good_lazy)
 # Works on JAX and Dask
 lazy_xp_function(non_materializable2, jax_jit=False, allow_dask_compute=2)
+lazy_xp_function(non_materializable3, jax_jit=False, allow_dask_compute=True)
 # Works on JAX, but not Dask
-lazy_xp_function(non_materializable3, jax_jit=False, allow_dask_compute=1)
+lazy_xp_function(non_materializable4, jax_jit=False, allow_dask_compute=1)
 # Works neither on Dask nor JAX
-lazy_xp_function(non_materializable4)
+lazy_xp_function(non_materializable5)
 
 
 def test_lazy_xp_function(xp: ModuleType):
@@ -217,29 +222,30 @@ def test_lazy_xp_function(xp: ModuleType):
     xp_assert_equal(non_materializable(x), xp.asarray([1.0, 2.0]))
     # Wrapping explicitly disabled
     xp_assert_equal(non_materializable2(x), xp.asarray([1.0, 2.0]))
+    xp_assert_equal(non_materializable3(x), xp.asarray([1.0, 2.0]))
 
     if is_jax_namespace(xp):
-        xp_assert_equal(non_materializable3(x), xp.asarray([1.0, 2.0]))
+        xp_assert_equal(non_materializable4(x), xp.asarray([1.0, 2.0]))
         with pytest.raises(
             TypeError, match="Attempted boolean conversion of traced array"
         ):
-            _ = non_materializable4(x)  # Wrapped
+            _ = non_materializable5(x)  # Wrapped
 
     elif is_dask_namespace(xp):
         with pytest.raises(
             AssertionError,
             match=r"dask\.compute.* 2 times, but only up to 1 calls are allowed",
         ):
-            _ = non_materializable3(x)
+            _ = non_materializable4(x)
         with pytest.raises(
             AssertionError,
             match=r"dask\.compute.* 1 times, but no calls are allowed",
         ):
-            _ = non_materializable4(x)
+            _ = non_materializable5(x)
 
     else:
-        xp_assert_equal(non_materializable3(x), xp.asarray([1.0, 2.0]))
         xp_assert_equal(non_materializable4(x), xp.asarray([1.0, 2.0]))
+        xp_assert_equal(non_materializable5(x), xp.asarray([1.0, 2.0]))
 
 
 def static_params(x: Array, n: int, flag: bool = False) -> Array: