Rename to apply_lazy; add as_numpy=False param

Author: crusaderky

docs/api-reference.md

@@ -6,7 +6,7 @@
     :nosignatures:
     :toctree: generated
 
-    apply_numpy_func
+    apply_lazy
     at
     atleast_nd
     cov

src/array_api_extra/__init__.py

@@ -1,7 +1,7 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._delegation import pad
-from ._lib._apply import apply_numpy_func
+from ._lib._apply import apply_lazy
 from ._lib._funcs import (
     at,
     atleast_nd,
@@ -19,7 +19,7 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
-    "apply_numpy_func",
+    "apply_lazy",
     "at",
     "atleast_nd",
     "cov",

src/array_api_extra/_lib/_apply.py

@@ -19,6 +19,7 @@
     from typing import ParamSpec, TypeAlias
 
     import numpy as np
+    from numpy.typing import ArrayLike
 
     NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[no-any-explicit]
     P = ParamSpec("P")
@@ -32,58 +33,74 @@ class P:  # pylint: disable=missing-class-docstring
 
 
 @overload
-def apply_numpy_func(  # type: ignore[valid-type]
-    func: Callable[P, NumPyObject],
+def apply_lazy(  # type: ignore[valid-type]
+    func: Callable[P, ArrayLike],
     *args: Array,
     shape: tuple[int | None, ...] | None = None,
     dtype: DType | None = None,
+    as_numpy: bool = False,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> Array: ...  # numpydoc ignore=GL08
 
 
 @overload
-def apply_numpy_func(  # type: ignore[valid-type]
-    func: Callable[P, Sequence[NumPyObject]],
+def apply_lazy(  # type: ignore[valid-type]
+    func: Callable[P, Sequence[ArrayLike]],
     *args: Array,
     shape: Sequence[tuple[int | None, ...]],
     dtype: Sequence[DType] | None = None,
+    as_numpy: bool = False,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
-def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
-    func: Callable[P, NumPyObject | Sequence[NumPyObject]],
+def apply_lazy(  # 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,
     dtype: DType | Sequence[DType] | None = None,
+    as_numpy: bool = False,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
 ) -> Array | tuple[Array, ...]:
     """
-    Apply a function that operates on NumPy arrays to Array API compliant arrays.
+    Lazily apply an eager function.
+
+    If the backend of the input arrays is lazy, e.g. Dask or jitted JAX, the execution
+    of the function is delayed until the graph is materialized; if it's eager, the
+    function is executed immediately.
 
     Parameters
     ----------
     func : callable
-        The function to apply. It must accept one or more NumPy arrays or generics as
-        positional arguments and return either a single NumPy array or generic, or a
-        tuple or list thereof.
+        The function to apply.
+
+        It must accept one or more array API compliant arrays as positional arguments.
+        If `as_numpy=True`, inputs are converted to NumPy before they are passed to
+        `func`.
+        It must return either a single array-like or a sequence of array-likes.
 
-        It must be a pure function, i.e. without side effects, as depending on the
+        `func` must be a pure function, i.e. without side effects, 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
-        :func:`numpy.asarray` to them to convert them to numpy; read notes below about
-        specific backends.
+        One or more Array API compliant arrays.
+
+        If `as_numpy=True`, you need to be able to apply :func:`numpy.asarray` to them
+        to convert them to numpy; read notes below about specific backends.
     shape : tuple[int | None, ...] | Sequence[tuple[int, ...]], optional
         Output shape or sequence of output shapes, one for each output of `func`.
         Default: assume single output and broadcast shapes of the input arrays.
     dtype : DType | Sequence[DType], optional
         Output dtype or sequence of output dtypes, one for each output of `func`.
         dtype(s) must belong to the same array namespace as the input arrays.
         Default: infer the result type(s) from the input arrays.
+    as_numpy : bool, optional
+        If True, convert the input arrays to NumPy before passing them to `func`.
+        This is particularly useful to make numpy-only functions, e.g. written in Cython
+        or Numba, work transparently API arrays.
+        Default: False.
     xp : array_namespace, optional
         The standard-compatible namespace for `args`. Default: infer.
     **kwargs : Any, optional
@@ -95,7 +112,7 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     Array | tuple[Array, ...]
         The result(s) of `func` applied to the input arrays, wrapped in the same
         array namespace as the inputs.
-        If shape is omitted or a `tuple[int, ...]`, this is a single array.
+        If shape is omitted or a `tuple[int | None, ...]`, this is a single array.
         Otherwise, it's a tuple of arrays.
 
     Notes
@@ -106,23 +123,26 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         When running inside `jax.jit`, `shape` must be fully known, i.e. it cannot
         contain any `None` elements.
 
-        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.
+        Using this with `as_numpy=False` is particularly useful to apply non-jittable
+        JAX functions to arrays on GPU devices.
+        If `as_numpy=True`, 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.
+        If `as_numpy=True`, 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
+        If `as_numpy=True`, 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; also, shape
+        `apply_lazy` 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.
 
@@ -136,7 +156,13 @@ def apply_numpy_func(  # 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_numpy_func`.
+        `apply_lazy`.
+
+    Dask wrapping around other backends
+        If `as_numpy=False`, `func` will receive in input eager arrays of the meta
+        namespace, as defined by the `._meta` attribute of the input Dask arrays.
+        The outputs of `func` will be wrapped by the meta namespace, and then wrapped
+        again by Dask.
 
     Raises
     ------
@@ -202,7 +228,10 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         metas = [arg._meta for arg in args if hasattr(arg, "_meta")]  # pylint: disable=protected-access
         meta_xp = array_namespace(*metas)
 
-        wrapped = dask.delayed(_npfunc_wrapper(func, multi_output, meta_xp), pure=True)
+        wrapped = dask.delayed(
+            _apply_lazy_wrapper(func, as_numpy, multi_output, meta_xp),
+            pure=True,
+        )
         # This finalizes each arg, which is the same as arg.rechunk(-1).
         # Please read docstring above for why we're not using
         # dask.array.map_blocks or dask.array.blockwise!
@@ -227,7 +256,7 @@ def apply_numpy_func(  # 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 = _npfunc_wrapper(func, multi_output, xp)
+        wrapped = _apply_lazy_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
@@ -251,19 +280,20 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
     else:
         # Eager backends
-        wrapped = _npfunc_wrapper(func, multi_output, xp)
+        wrapped = _apply_lazy_wrapper(func, as_numpy, multi_output, xp)
         out = wrapped(*args, **kwargs)
 
     return out if multi_output else out[0]
 
 
-def _npfunc_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01
-    func: Callable[..., NumPyObject | Sequence[NumPyObject]],
+def _apply_lazy_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_numpy_func`.
+    Helper of `apply_lazy`.
 
     Given a function that accepts one or more numpy arrays as positional arguments and
     returns a single numpy array or a sequence of numpy arrays, return a function that
@@ -284,19 +314,13 @@ def wrapper(  # type: ignore[no-any-decorated,no-any-explicit]
     ) -> tuple[Array, ...]:  # numpydoc ignore=GL08
         import numpy as np  # pylint: disable=import-outside-toplevel
 
-        args = tuple(np.asarray(arg) for arg in args)
+        if as_numpy:
+            args = tuple(np.asarray(arg) for arg in args)
         out = func(*args, **kwargs)
 
-        # Stay relaxed on output validation, e.g. in case func returns a
-        # Python scalar instead of a np.generic
         if multi_output:
-            if not isinstance(out, Sequence) or isinstance(out, np.ndarray):
-                msg = "Expected multiple outputs, got a single one"
-                raise ValueError(msg)
-            outs = out
-        else:
-            outs = [cast("NumPyObject", out)]
-
-        return tuple(xp.asarray(o) for o in outs)
+            assert isinstance(out, Sequence)
+            return tuple(xp.asarray(o) for o in out)
+        return (xp.asarray(out),)
 
     return wrapper