Use _implementation on all functions that have it in the array API su…

…bmodule That way they only work on actual ndarray inputs, not array-like, which is more inline with the spec.
numpy · charris · Aug 24, 2021 · Jan 11, 2021 · Jan 11, 2021 · Jan 11, 2021
commit 1efd55efa8cac9afd12d299dcf8912a7a7ac8a68
diff --git a/numpy/_array_api/_creation_functions.py b/numpy/_array_api/_creation_functions.py
@@ -35,7 +35,7 @@ def empty_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[d
     if device is not None:
         # Note: Device support is not yet implemented on ndarray
         raise NotImplementedError("Device support is not yet implemented")
-    return np.empty_like(x, dtype=dtype)
+    return np.empty_like._implementation(x, dtype=dtype)
 
 def eye(N: int, /, *, M: Optional[int] = None, k: Optional[int] = 0, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
@@ -68,7 +68,7 @@ def full_like(x: array, fill_value: Union[int, float], /, *, dtype: Optional[dty
     if device is not None:
         # Note: Device support is not yet implemented on ndarray
         raise NotImplementedError("Device support is not yet implemented")
-    return np.full_like(x, fill_value, dtype=dtype)
+    return np.full_like._implementation(x, fill_value, dtype=dtype)
 
 def linspace(start: Union[int, float], stop: Union[int, float], num: int, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None, endpoint: bool = True) -> array:
     """
@@ -101,7 +101,7 @@ def ones_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[de
     if device is not None:
         # Note: Device support is not yet implemented on ndarray
         raise NotImplementedError("Device support is not yet implemented")
-    return np.ones_like(x, dtype=dtype)
+    return np.ones_like._implementation(x, dtype=dtype)
 
 def zeros(shape: Union[int, Tuple[int, ...]], /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     """
@@ -123,4 +123,4 @@ def zeros_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[d
     if device is not None:
         # Note: Device support is not yet implemented on ndarray
         raise NotImplementedError("Device support is not yet implemented")
-    return np.zeros_like(x, dtype=dtype)
+    return np.zeros_like._implementation(x, dtype=dtype)
diff --git a/numpy/_array_api/_elementwise_functions.py b/numpy/_array_api/_elementwise_functions.py
@@ -381,7 +381,7 @@ def round(x: array, /) -> array:
 
     See its docstring for more information.
     """
-    return np.round(x)
+    return np.round._implementation(x)
 
 def sign(x: array, /) -> array:
     """

diff --git a/numpy/_array_api/_linear_algebra_functions.py b/numpy/_array_api/_linear_algebra_functions.py
@@ -18,7 +18,7 @@ def cross(x1: array, x2: array, /, *, axis: int = -1) -> array:
 
     See its docstring for more information.
     """
-    return np.cross(x1, x2, axis=axis)
+    return np.cross._implementation(x1, x2, axis=axis)
 
 def det(x: array, /) -> array:
     """
@@ -35,7 +35,7 @@ def diagonal(x: array, /, *, axis1: int = 0, axis2: int = 1, offset: int = 0) ->
 
     See its docstring for more information.
     """
-    return np.diagonal(x, axis1=axis1, axis2=axis2, offset=offset)
+    return np.diagonal._implementation(x, axis1=axis1, axis2=axis2, offset=offset)
 
 # def dot():
 #     """
@@ -128,7 +128,7 @@ def outer(x1: array, x2: array, /) -> array:
 
     See its docstring for more information.
     """
-    return np.outer(x1, x2)
+    return np.outer._implementation(x1, x2)
 
 # def pinv():
 #     """
@@ -176,12 +176,12 @@ def trace(x: array, /, *, axis1: int = 0, axis2: int = 1, offset: int = 0) -> ar
 
     See its docstring for more information.
     """
-    return np.asarray(np.trace(x, axis1=axis1, axis2=axis2, offset=offset))
+    return np.asarray(np.trace._implementation(x, axis1=axis1, axis2=axis2, offset=offset))
 
 def transpose(x: array, /, *, axes: Optional[Tuple[int, ...]] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.transpose <numpy.transpose>`.
 
     See its docstring for more information.
     """
-    return np.transpose(x, axes=axes)
+    return np.transpose._implementation(x, axes=axes)
diff --git a/numpy/_array_api/_manipulation_functions.py b/numpy/_array_api/_manipulation_functions.py
@@ -19,44 +19,44 @@ def expand_dims(x: array, axis: int, /) -> array:
 
     See its docstring for more information.
     """
-    return np.expand_dims(x, axis)
+    return np.expand_dims._implementation(x, axis)
 
 def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.flip <numpy.flip>`.
 
     See its docstring for more information.
     """
-    return np.flip(x, axis=axis)
+    return np.flip._implementation(x, axis=axis)
 
 def reshape(x: array, shape: Tuple[int, ...], /) -> array:
     """
     Array API compatible wrapper for :py:func:`np.reshape <numpy.reshape>`.
 
     See its docstring for more information.
     """
-    return np.reshape(x, shape)
+    return np.reshape._implementation(x, shape)
 
 def roll(x: array, shift: Union[int, Tuple[int, ...]], /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.roll <numpy.roll>`.
 
     See its docstring for more information.
     """
-    return np.roll(x, shift, axis=axis)
+    return np.roll._implementation(x, shift, axis=axis)
 
 def squeeze(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
     """
     Array API compatible wrapper for :py:func:`np.squeeze <numpy.squeeze>`.
 
     See its docstring for more information.
     """
-    return np.squeeze(x, axis=axis)
+    return np.squeeze._implementation(x, axis=axis)
 
 def stack(arrays: Tuple[array], /, *, axis: int = 0) -> array:
     """
     Array API compatible wrapper for :py:func:`np.stack <numpy.stack>`.
 
     See its docstring for more information.
     """
-    return np.stack(arrays, axis=axis)
+    return np.stack._implementation(arrays, axis=axis)
diff --git a/numpy/_array_api/_searching_functions.py b/numpy/_array_api/_searching_functions.py
@@ -11,7 +11,7 @@ def argmax(x: array, /, *, axis: int = None, keepdims: bool = False) -> array:
     See its docstring for more information.
     """
     # Note: this currently 
10000
fails as np.argmax does not implement keepdims
-    return np.asarray(np.argmax(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.argmax._implementation(x, axis=axis, keepdims=keepdims))
 
 def argmin(x: array, /, *, axis: int = None, keepdims: bool = False) -> array:
     """
@@ -20,20 +20,20 @@ def argmin(x: array, /, *, axis: int = None, keepdims: bool = False) -> array:
     See its docstring for more information.
     """
     # Note: this currently fails as np.argmin does not implement keepdims
-    return np.asarray(np.argmin(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.argmin._implementation(x, axis=axis, keepdims=keepdims))
 
 def nonzero(x: array, /) -> Tuple[array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.nonzero <numpy.nonzero>`.
 
     See its docstring for more information.
     """
-    return np.nonzero(x)
+    return np.nonzero._implementation(x)
 
 def where(condition: array, x1: array, x2: array, /) -> array:
     """
     Array API compatible wrapper for :py:func:`np.where <numpy.where>`.
 
     See its docstring for more information.
     """
-    return np.where(condition, x1, x2)
+    return np.where._implementation(condition, x1, x2)
diff --git a/numpy/_array_api/_set_functions.py b/numpy/_array_api/_set_functions.py
@@ -10,4 +10,4 @@ def unique(x: array, /, *, return_counts: bool = False, return_index: bool = Fal
 
     See its docstring for more information.
     """
-    return np.unique(x, return_counts=return_counts, return_index=return_index, return_inverse=return_inverse, sorted=sorted)
+    return np.unique._implementation(x, return_counts=return_counts, return_index=return_index, return_inverse=return_inverse, sorted=sorted)
diff --git a/numpy/_array_api/_sorting_functions.py b/numpy/_array_api/_sorting_functions.py
@@ -12,7 +12,7 @@ def argsort(x: array, /, *, axis: int = -1, descending: bool = False, stable: bo
     """
     # Note: this keyword argument is different, and the default is different.
     kind = 'stable' if stable else 'quicksort'
-    res = np.argsort(x, axis=axis, kind=kind)
+    res = np.argsort._implementation(x, axis=axis, kind=kind)
     if descending:
         res = np.flip(res, axis=axis)
     return res
@@ -25,7 +25,7 @@ def sort(x: array, /, *, axis: int = -1, descending: bool = False, stable: bool
     """
     # Note: this keyword argument is different, and the default is different.
     kind = 'stable' if stable else 'quicksort'
-    res = np.sort(x, axis=axis, kind=kind)
+    res = np.sort._implementation(x, axis=axis, kind=kind)
        if descending:
         res = np.flip(res, axis=axis)
     return res
diff --git a/numpy/_array_api/_statistical_functions.py b/numpy/_array_api/_statistical_functions.py
@@ -5,24 +5,24 @@
 import numpy as np
 
 def max(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
-def max(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
+from numpy.typing import _Shape  # i.e. `Tuple[int, ...]
+
+def max(x: array, /, *, axis: None | int | _Shape = None, keepdims: bool = False) -> array:
-def max(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
+from numpy.typing import _Shape  # i.e. `Tuple[int, ...]
+
+def max(x: array, /, *, axis: None | int | _Shape = None, keepdims: bool = False) -> array:
-    return np.max(x, axis=axis, keepdims=keepdims)
+    return np.max._implementation(x, axis=axis, keepdims=keepdims)
 
 def mean(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
-    return np.asarray(np.mean(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.mean._implementation(x, axis=axis, keepdims=keepdims))
 
 def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
-    return np.min(x, axis=axis, keepdims=keepdims)
+    return np.min._implementation(x, axis=axis, keepdims=keepdims)
 
 def prod(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
-    return np.asarray(np.prod(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.prod._implementation(x, axis=axis, keepdims=keepdims))
 
 def std(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, keepdims: bool = False) -> array:
     # Note: the keyword argument correction is different here
-    return np.asarray(np.std(x, axis=axis, ddof=correction, keepdims=keepdims))
+    return np.asarray(np.std._implementation(x, axis=axis, ddof=correction, keepdims=keepdims))
 
 def sum(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
-    return np.asarray(np.sum(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.sum._implementation(x, axis=axis, keepdims=keepdims))
 
 def var(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, keepdims: bool = False) -> array:
     # Note: the keyword argument correction is different here
-    return np.asarray(np.var(x, axis=axis, ddof=correction, keepdims=keepdims))
+    return np.asarray(np.var._implementation(x, axis=axis, ddof=correction, keepdims=keepdims))
diff --git a/numpy/_array_api/_utility_functions.py b/numpy/_array_api/_utility_functions.py
@@ -10,12 +10,12 @@ def all(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
 
     See its docstring for more information.
     """
-    return np.asarray(np.all(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.all._implementation(x, axis=axis, keepdims=keepdims))
 
 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>`.
 
     See its docstring for more information.
     """
-    return np.asarray(np.any(x, axis=axis, keepdims=keepdims))
+    return np.asarray(np.any._implementation(x, axis=axis, keepdims=keepdims))