Add multi-device test for take

Author: betatim

array_api_strict/_creation_functions.py

@@ -37,7 +37,7 @@ def _check_device(device):
     if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")
 
-    if device not in ALL_DEVICES:
+    if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
 
 def asarray(

array_api_strict/_indexing_functions.py

@@ -22,4 +22,6 @@ def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
         raise TypeError("Only integer dtypes are allowed in indexing")
     if indices.ndim != 1:
         raise ValueError("Only 1-dim indices array is supported")
-    return Array._new(np.take(x._array, indices._array, axis=axis))
+    if x.device != indices.device:
+        raise RuntimeError(f"Arrays from two different devices ({x.device} and {indices.device}) can not be combined.")
+    return Array._new(np.take(x._array, indices._array, axis=axis), device=x.device)

array_api_strict/tests/test_indexing_functions.py

@@ -22,3 +22,23 @@ def test_take_function(x, indices, axis, expected):
     indices = xp.asarray(indices)
     out = xp.take(x, indices, axis=axis)
     assert xp.all(out == xp.asarray(expected))
+
+
+def test_take_device():
+    x = xp.asarray([2, 3])
+    indices = xp.asarray([1, 1, 0])
+    xp.take(x, indices)
+
+    x = xp.asarray([2, 3])
+    indices = xp.asarray([1, 1, 0], device=xp.Device("device1"))
+    with pytest.raises(RuntimeError, match="Arrays from two different devices"):
+        xp.take(x, indices)
+
+    x = xp.asarray([2, 3], device=xp.Device("device1"))
+    indices = xp.asarray([1, 1, 0])
+    with pytest.raises(RuntimeError, match="Arrays from two different devices"):
+        xp.take(x, indices)
+
+    x = xp.asarray([2, 3], device=xp.Device("device1"))
+    indices = xp.asarray([1, 1, 0], device=xp.Device("device1"))
+    xp.take(x, indices)