DNM WIP ENH: More lazy functions

docs/api-lazy.md

@@ -0,0 +1,18 @@
+# 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
+    lazy_raise
+    lazy_wait_on
+    lazy_warn
+    testing.lazy_xp_function
+    testing.patch_lazy_xp_functions
+```

docs/conf.py

@@ -53,7 +53,10 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable", None),
+    "dask": ("https://docs.dask.org/en/stable", None),
     "jax": ("https://jax.readthedocs.io/en/latest", None),
+    "equinox": ("https://docs.kidger.site/equinox/", None),
 }
 
 nitpick_ignore = [

docs/index.md

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

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
     "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["array-api-compat>=1.10.0,<2"]
+# dependencies = ["array-api-compat>=1.10.0,<2"]  # DNM
 
 [project.urls]
 Homepage = "https://github.com/data-apis/array-api-extra"
@@ -48,10 +48,11 @@ 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.10.0,<2"  # DNM
 
 [tool.pixi.pypi-dependencies]
 array-api-extra = { path = ".", editable = true }
+array-api-compat = { git = "https://github.com/data-apis/array-api-compat" }  # DNM
 
 [tool.pixi.feature.lint.dependencies]
 typing-extensions = "*"
@@ -106,6 +107,7 @@ sphinx-autodoc-typehints = "*"
 dask-core = "*"
 pytest = "*"
 typing-extensions = "*"
+numpy = "*"
 
 [tool.pixi.feature.docs.tasks]
 docs = { cmd = "sphinx-build . build/", cwd = "docs" }
@@ -311,10 +313,5 @@ checks = [
     "ES01", # most docstrings do not need an extended summary
 ]
 exclude = [ # don't report on objects that match any of these regex
-    '.*test_at.*',
-    '.*test_funcs.*',
-    '.*test_testing.*',
-    '.*test_utils.*',
-    '.*test_version.*',
-    '.*test_vendor.*',
+    '.*test_*',
 ]

src/array_api_extra/__init__.py

@@ -12,6 +12,7 @@
     setdiff1d,
     sinc,
 )
+from ._lib._lazy import lazy_apply, lazy_raise, lazy_wait_on, lazy_warn
 
 __version__ = "0.6.1.dev0"
 
@@ -25,6 +26,10 @@
     "expand_dims",
     "isclose",
     "kron",
+    "lazy_apply",
+    "lazy_raise",
+    "lazy_wait_on",
+    "lazy_warn",
     "nunique",
     "pad",
     "setdiff1d",