Add type annotations to the array api submodule function definitions

Some stubs still need to be modified to properly pass mypy type checking.
Also, 'device' is just left as a TypeVar() for now.

1 files changed, 6 insertions, 2 deletions

diff --git a/numpy/_array_api/_utility_functions.py b/numpy/_array_api/_utility_functions.py
index 19743d15c..69e17e0e5 100644
--- a/numpy/_array_api/_utility_functions.py
+++ b/numpy/_array_api/_utility_functions.py
@@ -1,6 +1,10 @@
+from __future__ import annotations
+
+from ._types import Optional, Tuple, Union, array
+
 import numpy as np
 
-def all(x, /, *, axis=None, keepdims=False):
+def all(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
     """
     Array API compatible wrapper for :py:func:`np.all <numpy.all>`.
 
@@ -8,7 +12,7 @@ def all(x, /, *, axis=None, keepdims=False):
     """
     return np.all(x, axis=axis, keepdims=keepdims)
 
-def any(x, /, *, axis=None, keepdims=False):
+def any(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
     """
     Array API compatible wrapper for :py:func:`np.any <numpy.any>`.