ENH: lazy_apply (data-apis#86)

* ENH: New function `lazy_apply`

* Update docs/api-lazy.md

Co-authored-by: Ralf Gommers <[email protected]>

* Update src/array_api_extra/_lib/_lazy.py

Co-authored-by: Ralf Gommers <[email protected]>

* Code review

* Remove kwargs introspection; support None | complex args

* Don't always import numpy

* update lockfile

* appease mypy

---------

Co-authored-by: Ralf Gommers <[email protected]>
Co-authored-by: Lucas Colley <[email protected]>

Author: 3 people

docs/api-lazy.md

@@ -0,0 +1,15 @@
+# Tools for lazy backends
+
+These additional functions are meant to be used to support compatibility with
+lazy backends, e.g. Dask or JAX:
+
+```{eval-rst}
+.. currentmodule:: array_api_extra
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    lazy_apply
+    testing.lazy_xp_function
+    testing.patch_lazy_xp_functions
+```

docs/index.md

@@ -5,7 +5,7 @@
 :hidden:
 self
 api-reference.md
-testing-utils.md
+api-lazy.md
 contributing.md
 contributors.md
 ```

docs/testing-utils.md

@@ -108,6 +108,7 @@ sphinx-autodoc-typehints = "*"
 dask-core = "*"
 pytest = "*"
 typing-extensions = "*"
+numpy = "*"
 
 [tool.pixi.feature.docs.tasks]
 docs = { cmd = "sphinx-build . build/", cwd = "docs" }

pixi.lock

@@ -14,6 +14,7 @@
     setdiff1d,
     sinc,
 )
+from ._lib._lazy import lazy_apply
 
 __version__ = "0.7.0.dev0"
 
@@ -29,6 +30,7 @@
     "expand_dims",
     "isclose",
     "kron",
+    "lazy_apply",
     "nunique",
     "pad",
     "setdiff1d",

pyproject.toml

@@ -36,7 +36,7 @@
 
 
 @overload
-def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
+def apply_where(  # type: ignore[explicit-any,decorated-any] # numpydoc ignore=GL08
     cond: Array,
     args: Array | tuple[Array, ...],
     f1: Callable[..., Array],
@@ -48,7 +48,7 @@ def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ig
 
 
 @overload
-def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
+def apply_where(  # type: ignore[explicit-any,decorated-any] # numpydoc ignore=GL08
     cond: Array,
     args: Array | tuple[Array, ...],
     f1: Callable[..., Array],
@@ -59,7 +59,7 @@ def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ig
 ) -> Array: ...
 
 
-def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
+def apply_where(  # type: ignore[explicit-any] # numpydoc ignore=PR01,PR02
     cond: Array,
     args: Array | tuple[Array, ...],
     f1: Callable[..., Array],
@@ -145,7 +145,7 @@ def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
     return _apply_where(cond, f1, f2, fill_value, *args_, xp=xp)
 
 
-def _apply_where(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01
+def _apply_where(  # type: ignore[explicit-any]  # numpydoc ignore=PR01,RT01
     cond: Array,
     f1: Callable[..., Array],
     f2: Callable[..., Array] | None,
@@ -743,7 +743,7 @@ def pad(
         pad_width_seq = cast(list[tuple[int, int]], list(pad_width))
 
     # https://github.com/python/typeshed/issues/13376
-    slices: list[slice] = []  # type: ignore[no-any-explicit]
+    slices: list[slice] = []  # type: ignore[explicit-any]
     newshape: list[int] = []
     for ax, w_tpl in enumerate(pad_width_seq):
         if len(w_tpl) != 2: