nits

Author: crusaderky

src/array_api_extra/_lib/_apply.py

@@ -8,14 +8,12 @@
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, cast, overload
 
-from ._utils._compat import (
-    array_namespace,
-    is_dask_namespace,
-    is_jax_namespace
-)
+from ._utils._compat import array_namespace, is_dask_namespace, is_jax_namespace
 from ._utils._typing import Array, DType
 
 if TYPE_CHECKING:
+    # TODO move outside TYPE_CHECKING
+    # depends on scikit-learn abandoning Python 3.9
     # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
     from typing import ParamSpec, TypeAlias
 
@@ -72,8 +70,8 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         positional arguments and return either a single NumPy array or generic, or a
         tuple or list thereof.
 
-        It must be a pure function, i.e. without side effects such as disk output,
-        as depending on the backend it may be executed more than once.
+        It 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
@@ -225,18 +223,6 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         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}"
-                raise ValueError(msg)
-
     return out if multi_output else out[0]
 
 

src/array_api_extra/_lib/_utils/_compat.pyi

@@ -23,8 +23,8 @@ def is_cupy_namespace(xp: ModuleType, /) -> bool: ...
 def is_dask_namespace(xp: ModuleType, /) -> bool: ...
 def is_jax_namespace(xp: ModuleType, /) -> bool: ...
 def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
+def is_pydata_sparse_namespace(xp: ModuleType, /) -> bool: ...
 def is_torch_namespace(xp: ModuleType, /) -> bool: ...
 def is_jax_array(x: object, /) -> bool: ...
-def is_pydata_sparse_namespace(xp: ModuleType, /) -> bool: ...
 def is_writeable_array(x: object, /) -> bool: ...
 def size(x: Array, /) -> int | None: ...