apply_lazy -> lazy_apply

Author: crusaderky

docs/api-reference.md

@@ -6,13 +6,13 @@
     :nosignatures:
     :toctree: generated
 
-    apply_lazy
     at
     atleast_nd
     cov
     create_diagonal
     expand_dims
     kron
+    lazy_apply
     nunique
     pad
     setdiff1d

src/array_api_extra/__init__.py

@@ -1,7 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._delegation import pad
-from ._lib._apply import apply_lazy
 from ._lib._funcs import (
     at,
     atleast_nd,
@@ -13,19 +12,20 @@
     setdiff1d,
     sinc,
 )
+from ._lib._lazy import lazy_apply
 
 __version__ = "0.6.1.dev0"
 
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
-    "apply_lazy",
     "at",
     "atleast_nd",
     "cov",
     "create_diagonal",
     "expand_dims",
     "kron",
+    "lazy_apply",
     "nunique",
     "pad",
     "setdiff1d",

src/array_api_extra/_lib/_apply.py renamed to src/array_api_extra/_lib/_lazy.py

@@ -33,7 +33,7 @@ class P:  # pylint: disable=missing-class-docstring
 
 
 @overload
-def apply_lazy(  # type: ignore[valid-type]
+def lazy_apply(  # type: ignore[valid-type]
     func: Callable[P, ArrayLike],
     *args: Array,
     shape: tuple[int | None, ...] | None = None,
@@ -45,7 +45,7 @@ def apply_lazy(  # type: ignore[valid-type]
 
 
 @overload
-def apply_lazy(  # type: ignore[valid-type]
+def lazy_apply(  # type: ignore[valid-type]
     func: Callable[P, Sequence[ArrayLike]],
     *args: Array,
     shape: Sequence[tuple[int | None, ...]],
@@ -56,7 +56,7 @@ def apply_lazy(  # type: ignore[valid-type]
 ) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
-def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
+def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     func: Callable[P, Array | Sequence[ArrayLike]],
     *args: Array,
     shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None,
@@ -142,7 +142,7 @@ def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         This allows applying eager functions to dask arrays.
         The dask graph won't be computed.
 
-        `apply_lazy` doesn't know if `func` reduces along any axes; also, shape
+        `lazy_apply` 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.
 
@@ -156,7 +156,7 @@ def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         If you want to distribute the calculation across multiple workers, you
         should use :func:`dask.array.map_blocks`, :func:`dask.array.map_overlap`,
         :func:`dask.array.blockwise`, or a native Dask wrapper instead of
-        `apply_lazy`.
+        `lazy_apply`.
 
     Dask wrapping around other backends
         If `as_numpy=False`, `func` will receive in input eager arrays of the meta
@@ -229,7 +229,7 @@ def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         meta_xp = array_namespace(*metas)
 
         wrapped = dask.delayed(
-            _apply_lazy_wrapper(func, as_numpy, multi_output, meta_xp),
+            _lazy_apply_wrapper(func, as_numpy, multi_output, meta_xp),
             pure=True,
         )
         # This finalizes each arg, which is the same as arg.rechunk(-1).
@@ -256,7 +256,7 @@ def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
         import jax  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
 
-        wrapped = _apply_lazy_wrapper(func, as_numpy, multi_output, xp)
+        wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
 
         if any(s is None for shape in shapes for s in shape):
             # Unknown output shape. Won't work with jax.jit, but it
@@ -280,20 +280,20 @@ def apply_lazy(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
     else:
         # Eager backends
-        wrapped = _apply_lazy_wrapper(func, as_numpy, multi_output, xp)
+        wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
         out = wrapped(*args, **kwargs)
 
     return out if multi_output else out[0]
 
 
-def _apply_lazy_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01
+def _lazy_apply_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01
     func: Callable[..., ArrayLike | Sequence[ArrayLike]],
     as_numpy: bool,
     multi_output: bool,
     xp: ModuleType,
 ) -> Callable[..., tuple[Array, ...]]:
     """
-    Helper of `apply_lazy`.
+    Helper of `lazy_apply`.
 
     Given a function that accepts one or more arrays as positional arguments and returns
     a single array-like or a sequence of array-likes, return a function that accepts the