overload

Author: crusaderky

src/array_api_extra/_apply.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable, Sequence
 from functools import wraps
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, cast, overload
 
 from ._lib._compat import (
     array_namespace,
@@ -22,16 +22,39 @@
     import numpy as np
 
     NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[no-any-explicit]
+    KwArg: TypeAlias = Any  # type: ignore[no-any-explicit]
+
+
+@overload
+def apply_numpy_func(
+    func: Callable[..., NumPyObject],
+    *args: Array,
+    shape: tuple[int, ...] | None = None,
+    dtype: DType | None = None,
+    xp: ModuleType | None = None,
+    **kwargs: KwArg,
+) -> Array: ...  # numpydoc ignore=GL08
+
+
+@overload
+def apply_numpy_func(  # type: ignore[no-any-decorated]
+    func: Callable[..., Sequence[NumPyObject]],
+    *args: Array,
+    shape: Sequence[tuple[int, ...]],
+    dtype: Sequence[DType] | None = None,
+    xp: ModuleType | None = None,
+    **kwargs: Any,
+) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
 
 
 def apply_numpy_func(  # type: ignore[no-any-explicit]
     func: Callable[..., NumPyObject | Sequence[NumPyObject]],
     *args: Array,
-    shapes: Sequence[tuple[int, ...]] | None = None,
-    dtypes: Sequence[DType] | None = None,
+    shape: tuple[int, ...] | Sequence[tuple[int, ...]] | None = None,
+    dtype: DType | Sequence[DType] | None = None,
     xp: ModuleType | None = None,
     **kwargs: Any,
-) -> tuple[Array, ...]:
+) -> Array | tuple[Array, ...]:
     """
     Apply a function that operates on NumPy arrays to Array API compliant arrays.
 
@@ -48,15 +71,11 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
         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
         specific backends.
-    shapes : Sequence[tuple[int, ...]], optional
-        Sequence of output shapes, one for each output of `func`.
-        If `func` returns a single (non-sequence) output, this must be a sequence
-        with a single element.
-        Default: assume a single output and broadcast shapes of the input arrays.
-    dtypes : Sequence[DType], optional
-        Sequence of output dtypes, one for each output of `func`.
-        If `func` returns a single (non-sequence) output, this must be a sequence
-        with a single element.
+    shape : tuple[int, ...] | 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`.
         Default: infer the result type(s) from the input arrays.
     xp : array_namespace, optional
         The standard-compatible namespace for `args`. Default: infer.
@@ -66,9 +85,11 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
 
     Returns
     -------
-    tuple[Array, ...]
-        The result(s) of `func` applied to the input arrays.
-        This is always a tuple, even if `func` returns a single output.
+    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.
+        Otherwise, it's a tuple of arrays.
 
     Notes
     -----
@@ -110,46 +131,67 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
     """
     if xp is None:
         xp = array_namespace(*args)
-    if shapes is None:
+
+    # Normalize and validate shape and dtype
+    multi_output = False
+    if shape is None:
         shapes = [xp.broadcast_shapes(*(arg.shape for arg in args))]
-    if dtypes is None:
+    elif isinstance(shape, tuple) and all(isinstance(s, int) for s in shape):
+        shapes = [shape]
+    else:
+        shapes = shape
+        multi_output = True
+
+    if dtype is None:
         dtypes = [xp.result_type(*args)] * len(shapes)
+    elif multi_output:
+        if not isinstance(dtype, Sequence):
+            msg = "Got sequence of shapes but only one dtype"
+            raise TypeError(msg)
+        dtypes = dtype
+    else:
+        if isinstance(dtype, Sequence):
+            msg = "Got single shape but multiple dtypes"
+            raise TypeError(msg)
+        dtypes = [dtype]
 
     if len(shapes) != len(dtypes):
-        msg = f"got {len(shapes)} shapes and {len(dtypes)} dtypes"
+        msg = f"Got {len(shapes)} shapes and {len(dtypes)} dtypes"
         raise ValueError(msg)
     if len(shapes) == 0:
-        msg = "Must have at least one output array"
+        msg = "func must return one or more output arrays"
         raise ValueError(msg)
+    del shape
+    del dtype
 
+    # Backend-specific branches
     if is_dask_namespace(xp):
         import dask  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
 
         metas = [arg._meta for arg in args if hasattr(arg, "_meta")]  # pylint: disable=protected-access
         meta_xp = array_namespace(*metas)
-        meta = metas[0]
 
-        wrapped = dask.delayed(_npfunc_wrapper(func, meta_xp), pure=True)
+        wrapped = dask.delayed(_npfunc_wrapper(func, 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!
         delayed_out = wrapped(*args, **kwargs)
 
-        return tuple(
-            xp.from_delayed(delayed_out[i], shape=shape, dtype=dtype, meta=meta)
+        out = tuple(
+            xp.from_delayed(delayed_out[i], shape=shape, dtype=dtype, meta=metas[0])
             for i, (shape, dtype) in enumerate(zip(shapes, dtypes, strict=True))
         )
 
-    wrapped = _npfunc_wrapper(func, xp)
-    if is_jax_namespace(xp):
+    elif is_jax_namespace(xp):
         # If we're inside jax.jit, we can't eagerly convert
         # the JAX tracer objects to numpy.
         # Instead, we delay calling wrapped, which will receive
         # as arguments and will return JAX eager arrays.
 
         import jax  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
 
-        return cast(
+        wrapped = _npfunc_wrapper(func, multi_output, xp)
+        out = cast(
             tuple[Array, ...],
             jax.pure_callback(
                 wrapped,
@@ -162,25 +204,29 @@ def apply_numpy_func(  # type: ignore[no-any-explicit]
             ),
         )
 
-    # Eager backends
-    out = wrapped(*args, **kwargs)
+    else:
+        # Eager backends
+        wrapped = _npfunc_wrapper(func, multi_output, xp)
+        out = wrapped(*args, **kwargs)
 
-    # Output validation
-    if len(out) != len(shapes):
-        msg = f"func was declared to return {len(shapes)} outputs, got {len(out)}"
-        raise ValueError(msg)
-    for out_i, shape_i, dtype_i in zip(out, shapes, dtypes, strict=True):
-        if out_i.shape != shape_i:
-            msg = f"expected shape {shape_i}, got {out_i.shape}"
-            raise ValueError(msg)
-        if not xp.isdtype(out_i.dtype, dtype_i):
-            msg = f"expected dtype {dtype_i}, got {out_i.dtype}"
+        # Output validation
+        if len(out) != len(shapes):
+            msg = f"func was declared to return {len(shapes)} outputs, got {len(out)}"
             raise ValueError(msg)
-    return out  # type: ignore[no-any-return]
+        for out_i, shape_i, dtype_i in zip(out, shapes, dtypes, strict=True):
+            if out_i.shape != shape_i:
+                msg = f"expected shape {shape_i}, got {out_i.shape}"
+                raise ValueError(msg)
+            if not xp.isdtype(out_i.dtype, dtype_i):
+                msg = f"expected dtype {dtype_i}, got {out_i.dtype}"
+                raise ValueError(msg)
+
+    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]],
+    multi_output: bool,
     xp: ModuleType,
 ) -> Callable[..., tuple[Array, ...]]:
     """
@@ -208,14 +254,12 @@ def wrapper(  # type: ignore[no-any-decorated,no-any-explicit]
         args = tuple(np.asarray(arg) for arg in args)
         out = func(*args, **kwargs)
 
-        if isinstance(out, np.ndarray | 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)
+        else:
             out = (out,)
-        elif not isinstance(out, Sequence):  # pyright: ignore[reportUnnecessaryIsInstance]
-            msg = (
-                "apply_numpy_func: func must return a numpy object or a "
-                f"sequence of numpy objects; got {out}"
-            )
-            raise TypeError(msg)
 
         return tuple(xp.asarray(o) for o in out)