Add multi-device support for searching

Author: betatim

Diff for: array_api_strict/_searching_functions.py

@@ -19,7 +19,7 @@ def argmax(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     """
     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)))
+    return Array._new(np.asarray(np.argmax(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
 
 def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> Array:
@@ -30,7 +30,7 @@ def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     """
     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)))
+    return Array._new(np.asarray(np.argmin(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
 
 @requires_data_dependent_shapes
@@ -61,12 +61,16 @@ def searchsorted(
     """
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in searchsorted")
+
+    if x1.device != x2.device:
+        raise RuntimeError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
+
     sorter = sorter._array if sorter is not None else None
     # TODO: The sort order of nans and signed zeros is implementation
     # dependent. Should we error/warn if they are present?
 
     # x1 must be 1-D, but NumPy already requires this.
-    return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter))
+    return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter), device=x1.device)
 
 def where(condition: Array, x1: Array, x2: Array, /) -> Array:
     """
@@ -76,5 +80,9 @@ def where(condition: Array, x1: Array, x2: Array, /) -> Array:
     """
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
+
+    if len({a.device for a in (condition, x1, x2)}) > 1:
+        raise ValueError("where inputs must all be on the same device")
+
     x1, x2 = Array._normalize_two_args(x1, x2)
-    return Array._new(np.where(condition._array, x1._array, x2._array))
+    return Array._new(np.where(condition._array, x1._array, x2._array), device=x1.device)