Add multi-device support to sorting functions

Author: betatim

array_api_strict/_sorting_functions.py

@@ -33,7 +33,7 @@ def argsort(
         normalised_axis = axis if axis >= 0 else x.ndim + axis
         max_i = x.shape[normalised_axis] - 1
         res = max_i - res
-    return Array._new(res)
+    return Array._new(res, device=x.device)
 
 # Note: the descending keyword argument is new in this function
 def sort(
@@ -51,4 +51,4 @@ def sort(
     res = np.sort(x._array, axis=axis, kind=kind)
     if descending:
         res = np.flip(res, axis=axis)
-    return Array._new(res)
+    return Array._new(res, device=x.device)

array_api_strict/tests/test_sorting_functions.py

@@ -21,3 +21,17 @@ def test_stable_desc_argsort(obj, axis, expected):
     x = xp.asarray(obj)
     out = xp.argsort(x, axis=axis, stable=True, descending=True)
     assert xp.all(out == xp.asarray(expected))
+
+
+def test_argsort_device():
+    x = xp.asarray([1., 2., -1., 3.141], device=xp.Device("device1"))
+    y = xp.argsort(x)
+
+    assert y.device == x.device
+
+
+def test_sort_device():
+    x = xp.asarray([1., 2., -1., 3.141], device=xp.Device("device1"))
+    y = xp.sort(x)
+
+    assert y.device == x.device