docs

Author: crusaderky

docs/conf.py

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

pixi.lock

@@ -113,6 +113,7 @@ furo = ">=2023.08.17"
 myst-parser = ">=0.13"
 sphinx-copybutton = "*"
 sphinx-autodoc-typehints = "*"
+numpy = "*"
 
 [tool.pixi.feature.docs.tasks]
 docs = { cmd = "sphinx-build . build/", cwd = "docs" }

pyproject.toml

@@ -23,6 +23,13 @@
 
     NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[no-any-explicit]
     P = ParamSpec("P")
+else:
+    # Sphinx hacks
+    NumPyObject = Any
+
+    class P:  # pylint: disable=missing-class-docstring
+        args: tuple
+        kwargs: dict
 
 
 @overload
@@ -47,7 +54,7 @@ def apply_numpy_func(  # type: ignore[valid-type]
 ) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
-def apply_numpy_func(  # type: ignore[valid-type]
+def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     func: Callable[P, NumPyObject | Sequence[NumPyObject]],
     *args: Array,
     shape: tuple[int, ...] | Sequence[tuple[int, ...]] | None = None,
@@ -69,7 +76,7 @@ def apply_numpy_func(  # type: ignore[valid-type]
         as depending on the backend it may be executed more than once.
     *args : Array
         One or more Array API compliant arrays. You need to be able to apply
-        ``np.asarray()`` to them to convert them to numpy; read notes below about
+        :func:`numpy.asarray` to them to convert them to numpy; read notes below about
         specific backends.
     shape : tuple[int, ...] | Sequence[tuple[int, ...]], optional
         Output shape or sequence of output shapes, one for each output of `func`.
@@ -97,25 +104,23 @@ def apply_numpy_func(  # type: ignore[valid-type]
         This allows applying eager functions to jitted JAX arrays, which are lazy.
         The function won't be applied until the JAX array is materialized.
 
-        The `JAX transfer guard
-        <https://jax.readthedocs.io/en/latest/transfer_guard.html>`_
-        may prevent arrays on a GPU device from being transferred back to CPU.
-        This is treated as an implicit transfer.
+        The :doc:`jax:transfer_guard` may prevent arrays on a GPU device from being
+        transferred back to CPU. This is treated as an implicit transfer.
 
     PyTorch, CuPy
         These backends raise by default if you attempt to convert arrays on a GPU device
         to NumPy.
 
     Sparse
-        By default, sparse prevents implicit densification through ``np.asarray`.
-        `This safety mechanism can be disabled
+        By default, sparse prevents implicit densification through
+        :func:`numpy.asarray`. `This safety mechanism can be disabled
         <https://sparse.pydata.org/en/stable/operations.html#package-configuration>`_.
 
     Dask
         This allows applying eager functions to dask arrays.
         The dask graph won't be computed.
 
-        `apply_numpy_func` doesn't know if `func` reduces along any axes and shape
+        `apply_numpy_func` doesn't know if `func` reduces along any axes; also, shape
         changes are non-trivial in chunked Dask arrays. For these reasons, all inputs
         will be rechunked into a single chunk.
 
@@ -125,9 +130,19 @@ def apply_numpy_func(  # type: ignore[valid-type]
 
         The outputs will also be returned as a single chunk and you should consider
         rechunking them into smaller chunks afterwards.
+
         If you want to distribute the calculation across multiple workers, you
-        should use `dask.array.map_blocks`, `dask.array.blockwise`,
-        `dask.array.map_overlap`, or a native Dask wrapper instead of this function.
+        should use :func:`dask.array.map_blocks`, :func:`dask.array.map_overlap`,
+        :func:`dask.array.blockwise`, or a native Dask wrapper instead of
+        `apply_numpy_func`.
+
+    See Also
+    --------
+    jax.transfer_guard
+    jax.pure_callback
+    dask.array.map_blocks
+    dask.array.map_overlap
+    dask.array.blockwise
     """
     if xp is None:
         xp = array_namespace(*args)
@@ -239,8 +254,8 @@ def _npfunc_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT
 
     Any keyword arguments are passed through verbatim to the wrapped function.
 
-    Raise if np.asarray() raises on any input. This typically happens if the input is
-    lazy and has a guard against being implicitly turned into a NumPy array (e.g.
+    Raise if np.asarray raises on any input. This typically happens if the input is lazy
+    and has a guard against being implicitly turned into a NumPy array (e.g.
     densification for sparse arrays, device->host transfer for cupy and torch arrays).
     """