ENH Test tools for jax.jit and dask

data-apis · Jan 21, 2025 · 57e6cd0 · 57e6cd0
1 parent 6ee70c0
commit 57e6cd0
Show file tree

Hide file tree

Showing 11 changed files with 457 additions and 28 deletions.
diff --git a/docs/api-reference.md b/docs/api-reference.md
@@ -18,3 +18,15 @@
     setdiff1d
     sinc
 ```
+
+## Testing utilities
+
+```{eval-rst}
+.. currentmodule:: array_api_extra.testing
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    lazy_xp_function
+    patch_lazy_xp_functions
+```
diff --git a/pixi.lock b/pixi.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -98,6 +98,8 @@ furo = ">=2023.08.17"
 myst-parser = ">=0.13"
 sphinx-copybutton = "*"
 sphinx-autodoc-typehints = "*"
+# Needed to import parsed modules with autodoc
+pytest = "*"
 
 [tool.pixi.feature.docs.tasks]
 docs = { cmd = "sphinx-build . build/", cwd = "docs" }
@@ -180,8 +182,10 @@ markers = ["skip_xp_backend(library, *, reason=None): Skip test for a specific b
 
 [tool.coverage]
 run.source = ["array_api_extra"]
-report.exclude_also = ['\.\.\.']
-
+report.exclude_also = [
+    '\.\.\.',
+    'if TYPE_CHECKING:',
+]
 
 # mypy
 
@@ -221,6 +225,8 @@ reportMissingImports = false
 reportMissingTypeStubs = false
 # false positives for input validation
 reportUnreachable = false
+# ruff handles this
+reportUnusedParameter = false
 
 executionEnvironments = [
     { root = "tests", reportPrivateUsage = false },
@@ -282,7 +288,10 @@ messages_control.disable = [
     "design",                     # ignore heavily opinionated design checks
     "fixme",                      # allow FIXME comments
     "line-too-long",              # ruff handles this
+    "unused-argument",            # ruff handles this
     "missing-function-docstring", # numpydoc handles this
+    "import-error",               # mypy handles this
+    "import-outside-toplevel",    # optional dependencies
 ]
 
 

diff --git a/src/array_api_extra/_lib/_at.py b/src/array_api_extra/_lib/_at.py
@@ -188,7 +188,7 @@ def __getitem__(self, idx: Index, /) -> at:  # numpydoc ignore=PR01,RT01
     def _update_common(
         self,
         at_op: _AtOp,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -253,7 +253,7 @@ def _update_common(
 
     def set(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -269,8 +269,8 @@ def set(
     def _iop(
         self,
         at_op: _AtOp,
-        elwise_op: Callable[[Array, Array], Array],
-        y: Array,
+        elwise_op: Callable[[Array, Array | object], Array],
+        y: Array | object,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -294,7 +294,7 @@ def _iop(
 
     def add(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -308,7 +308,7 @@ def add(
 
     def subtract(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -318,7 +318,7 @@ def subtract(
 
     def multiply(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -328,7 +328,7 @@ def multiply(
 
     def divide(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -338,7 +338,7 @@ def divide(
 
     def power(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -348,7 +348,7 @@ def power(
 
     def min(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -361,7 +361,7 @@ def min(
 
     def max(
         self,
-        y: Array,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,

diff --git a/src/array_api_extra/_lib/_testing.py b/src/array_api_extra/_lib/_testing.py
@@ -2,6 +2,7 @@
 Testing utilities.
 
 Note that this is private API; don't expect it to be stable.
+See also ..testing for public testing utilities.
 """
 
 import math