DEV/MAINT: set up doctests and clean examples

Author: fbourgey

.github/workflows/ci.yml

@@ -38,6 +38,23 @@ jobs:
       - name: Lint (if this step fails, please 'pixi run lint' locally and push the changes)
         run: pixi run -e lint lint
 
+  doctests:
+    name: Doctests
+    runs-on: ubuntu-slim
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+
+      - uses: prefix-dev/setup-pixi@a09b6247153796b190642a2b53fac4241043cf6f # v0.10.0
+        with:
+          pixi-version: v0.71.2
+          cache: true
+          environments: tests
+
+      - name: Test public API examples
+        run: pixi run -e tests doctests
+
   checks:
     name: ${{ matrix.environment }} (${{ matrix.platform }})
     runs-on: >-

conftest.py

@@ -0,0 +1,22 @@
+"""Configure public API doctests."""
+
+import warnings
+from contextlib import contextmanager
+
+from scipy_doctest.conftest import dt_config
+
+
+@contextmanager
+def _doctest_context(_test=None):
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"`xpx\.(broadcast_shapes|expand_dims)` is deprecated.*",
+            category=DeprecationWarning,
+        )
+        yield
+
+
+dt_config.rtol = 1e-7
+dt_config.strict_check = True
+dt_config.user_context_mgr = _doctest_context

docs/contributing.md

@@ -16,6 +16,7 @@ All development tasks are then available via `pixi run`:
 ```bash
 pixi run tests      # run the tests
 pixi run open-docs  # build and preview the docs
+pixi run doctests   # run the doctests in the docs  
 pixi run lint       # run the full lint suite
 pixi run ipython    # spawn an ipython prompt with array-api-extra installed
 pixi run hooks      # install pre-commit hooks

pixi.lock

@@ -220,13 +220,19 @@ hypothesis = ">=6.155.7"
 array-api-strict = ">=2.6.1"
 numpy = ">=1.22.0"
 scipy = ">=1.15.2"
+scipy-doctest = ">=2,<3"
 
 [feature.tests.tasks]
 tests = {
   description = "Run tests",
   cmd = "pytest -v",
   default-environment = "tests",
 }
+doctests = {
+  description = "Run public API doctests",
+  cmd = "pytest --pyargs array_api_extra --doctest-modules --doctest-collect=api --doctest-only-doctests=true",
+  default-environment = "tests",
+}
 tests-cov = {
   description = "Run tests with coverage",
   cmd = "pytest -v -ra --cov --cov-report=xml --cov-report=term --durations=20",

pixi.toml

@@ -907,23 +907,23 @@ def nan_to_num(
     --------
     >>> import array_api_extra as xpx
     >>> import array_api_strict as xp
-    >>> xpx.nan_to_num(xp.inf)
-    1.7976931348623157e+308
-    >>> xpx.nan_to_num(-xp.inf)
-    -1.7976931348623157e+308
-    >>> xpx.nan_to_num(xp.nan)
-    0.0
+    >>> xpx.nan_to_num(xp.inf, xp=xp)
+    Array(1.79769313e+308, dtype=array_api_strict.float64)
+    >>> xpx.nan_to_num(-xp.inf, xp=xp)
+    Array(-1.79769313e+308, dtype=array_api_strict.float64)
+    >>> xpx.nan_to_num(xp.nan, xp=xp)
+    Array(0., dtype=array_api_strict.float64)
     >>> x = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128])
     >>> xpx.nan_to_num(x)
-    array([ 1.79769313e+308, -1.79769313e+308,  0.00000000e+000, # may vary
-           -1.28000000e+002,  1.28000000e+002])
+    Array([ 1.79769313e+308, -1.79769313e+308,  0.00000000e+000,
+           -1.28000000e+002,  1.28000000e+002],
+          dtype=array_api_strict.float64)
     >>> y = xp.asarray([complex(xp.inf, xp.nan), xp.nan, complex(xp.nan, xp.inf)])
-    array([  1.79769313e+308,  -1.79769313e+308,   0.00000000e+000, # may vary
-         -1.28000000e+002,   1.28000000e+002])
     >>> xpx.nan_to_num(y)
-    array([  1.79769313e+308 +0.00000000e+000j, # may vary
-             0.00000000e+000 +0.00000000e+000j,
-             0.00000000e+000 +1.79769313e+308j])
+    Array([1.79769313e+308+0.00000000e+000j,
+           0.00000000e+000+0.00000000e+000j,
+           0.00000000e+000+1.79769313e+308j],
+          dtype=array_api_strict.complex128)
     """
     if isinstance(fill_value, complex):
         msg = "Complex fill values are not supported."
@@ -1818,11 +1818,10 @@ def unravel_index(
     >>> import array_api_extra as xpx
     >>> import array_api_strict as xp
     >>> xs, ys = xpx.unravel_index(xp.asarray([1, 2, 4, 5, 6, 8]), (4, 3))
-    >>> xs, ys
-    (
-        Array([0, 0, 1, 1, 2, 2], dtype=array_api_strict.int64),
-        Array([1, 2, 1, 2, 0, 2], dtype=array_api_strict.int64),
-    )
+    >>> xs
+    Array([0, 0, 1, 1, 2, 2], dtype=array_api_strict.int64)
+    >>> ys
+    Array([1, 2, 1, 2, 0, 2], dtype=array_api_strict.int64)
     >>> [(int(x), int(y)) for x, y in zip(xs, ys)]
     [(0, 1), (0, 2), (1, 1), (1, 2), (2, 0), (2, 2)]
     >>> xs, ys = xpx.unravel_index(xp.arange(6), (2, 2))

src/array_api_extra/_delegation.py

@@ -76,9 +76,10 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
 
         You may use two alternate syntaxes::
 
-          >>> import array_api_extra as xpx
-          >>> xpx.at(x, idx).set(value)  # or add(value), etc.
-          >>> xpx.at(x)[idx].set(value)
+          import array_api_extra as xpx
+
+          xpx.at(x, idx).set(value)  # or add(value), etc.
+          xpx.at(x)[idx].set(value)
 
     copy : bool, optional
         None (default)
@@ -103,8 +104,9 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     (a) When you omit the ``copy`` parameter, you should never reuse the parameter
     array later on; ideally, you should reassign it immediately::
 
-        >>> import array_api_extra as xpx
-        >>> x = xpx.at(x, 0).set(2)
+        import array_api_extra as xpx
+
+        x = xpx.at(x, 0).set(2)
 
     The above best practice pattern ensures that the behaviour won't change depending
     on whether ``x`` is writeable or not, as the original ``x`` object is dereferenced
@@ -114,9 +116,9 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     On the reverse, the anti-pattern below must be avoided, as it will result in
     different behaviour on read-only versus writeable arrays::
 
-        >>> x = xp.asarray([0, 0, 0])
-        >>> y = xpx.at(x, 0).set(2)
-        >>> z = xpx.at(x, 1).set(3)
+        x = xp.asarray([0, 0, 0])
+        y = xpx.at(x, 0).set(2)
+        z = xpx.at(x, 1).set(3)
 
     In the above example, both calls to ``xpx.at`` update ``x`` in place *if possible*.
     This causes the behaviour to diverge depending on whether ``x`` is writeable or not:
@@ -129,22 +131,23 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
     The correct pattern to use if you want diverging outputs from the same input is
     to enforce copies::
 
-        >>> x = xp.asarray([0, 0, 0])
-        >>> y = xpx.at(x, 0).set(2, copy=True)  # Never updates x
-        >>> z = xpx.at(x, 1).set(3)  # May or may not update x in place
-        >>> del x  # avoid accidental reuse of x as we don't know its state anymore
+        x = xp.asarray([0, 0, 0])
+        y = xpx.at(x, 0).set(2, copy=True)  # Never updates x
+        z = xpx.at(x, 1).set(3)  # May or may not update x in place
+        del x  # avoid accidental reuse of x as we don't know its state anymore
 
     (b) The array API standard does not support integer array indices.
     The behaviour of update methods when the index is an array of integers is
     undefined and will vary between backends; this is particularly true when the
     index contains multiple occurrences of the same index, e.g.::
 
-        >>> import numpy as np
-        >>> import jax.numpy as jnp
-        >>> import array_api_extra as xpx
-        >>> xpx.at(np.asarray([123]), np.asarray([0, 0])).add(1)
+        import numpy as np
+        import jax.numpy as jnp
+        import array_api_extra as xpx
+
+        xpx.at(np.asarray([123]), np.asarray([0, 0])).add(1)
         array([124])
-        >>> xpx.at(jnp.asarray([123]), jnp.asarray([0, 0])).add(1)
+        xpx.at(jnp.asarray([123]), jnp.asarray([0, 0])).add(1)
         Array([125], dtype=int32)
 
     See Also
@@ -164,38 +167,39 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
 
     This pattern::
 
-        >>> mask = m(x)
-        >>> x[mask] = f(x[mask])
+        mask = m(x)
+        x[mask] = f(x[mask])
 
     Can't be replaced by `at`, as it won't work on Dask and JAX inside jax.jit::
 
-        >>> mask = m(x)
-        >>> x = xpx.at(x, mask).set(f(x[mask])  # Crash on Dask and jax.jit
+        mask = m(x)
+        x = xpx.at(x, mask).set(f(x[mask]))  # Crash on Dask and jax.jit
 
     You should instead use::
 
-        >>> x = xp.where(m(x), f(x), x)
+        x = xp.where(m(x), f(x), x)
 
     Examples
     --------
     Given either of these equivalent expressions::
 
-      >>> import array_api_extra as xpx
-      >>> x = xpx.at(x)[1].add(2)
-      >>> x = xpx.at(x, 1).add(2)
+        import array_api_extra as xpx
+
+        x = xpx.at(x)[1].add(2)
+        x = xpx.at(x, 1).add(2)
 
     If x is a JAX array, they are the same as::
 
-      >>> x = x.at[1].add(2)
+        x = x.at[1].add(2)
 
     If x is a read-only NumPy array, they are the same as::
 
-      >>> x = x.copy()
-      >>> x[1] += 2
+        x = x.copy()
+        x[1] += 2
 
     For other known backends, they are the same as::
 
-      >>> x[1] += 2
+        x[1] += 2
     """
 
     _x: Array

src/array_api_extra/_lib/_at.py

@@ -144,12 +144,12 @@ def apply_where(  # numpydoc ignore=PR01,PR02
     --------
     >>> import array_api_strict as xp
     >>> import array_api_extra as xpx
-    >>> a = xp.asarray([5, 4, 3])
-    >>> b = xp.asarray([0, 2, 2])
+    >>> a = xp.asarray([5.0, 4.0, 3.0])
+    >>> b = xp.asarray([0.0, 2.0, 2.0])
     >>> def f(a, b):
     ...     return a // b
     >>> xpx.apply_where(b != 0, (a, b), f, fill_value=xp.nan)
-    array([ nan,  2., 1.])
+    Array([nan,  2.,  1.], dtype=array_api_strict.float64)
     """
     # Parse and normalize arguments
     if (f2 is None) == (fill_value is None):

src/array_api_extra/_lib/_funcs.py

@@ -149,11 +149,11 @@ def lazy_xp_function(
 
         In other words, the pattern that is being tested is::
 
-            >>> @jax.jit
-            ... def user_func(x):
-            ...     y = user_prepares_inputs(x)
-            ...     z = func(y, some_static_arg=True)
-            ...     return user_consumes(z)
+            @jax.jit
+            def user_func(x):
+                y = user_prepares_inputs(x)
+                z = func(y, some_static_arg=True)
+                return user_consumes(z)
 
         Default: True.
     static_argnums : Deprecated