Update dtype strictness in array_api searching and sorting functions

Original NumPy Commit: 837b1af70ecea4877c8b1fee327d73d6dace517a

Author: asmeurer

array_api_strict/_searching_functions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from ._array_object import Array
-from ._dtypes import _result_type
+from ._dtypes import _result_type, _real_numeric_dtypes
 
 from typing import Optional, Tuple
 
@@ -14,6 +14,8 @@ def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
 
     See its docstring for more information.
     """
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in argmax")
     return Array._new(np.asarray(np.argmax(x._array, axis=axis, keepdims=keepdims)))
 
 
@@ -23,6 +25,8 @@ def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
 
     See its docstring for more information.
     """
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in argmin")
     return Array._new(np.asarray(np.argmin(x._array, axis=axis, keepdims=keepdims)))
 
 

array_api_strict/_sorting_functions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from ._array_object import Array
+from ._dtypes import _real_numeric_dtypes
 
 import numpy as np
 
@@ -14,6 +15,8 @@ def argsort(
 
     See its docstring for more information.
     """
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in argsort")
     # Note: this keyword argument is different, and the default is different.
     kind = "stable" if stable else "quicksort"
     if not descending:
@@ -41,6 +44,8 @@ def sort(
 
     See its docstring for more information.
     """
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in sort")
     # Note: this keyword argument is different, and the default is different.
     kind = "stable" if stable else "quicksort"
     res = np.sort(x._array, axis=axis, kind=kind)