add device check

Author: lucyleeow

array_api_strict/_array_object.py

@@ -727,7 +727,14 @@ def __getitem__(
         if isinstance(key, Array):
             key = (key,)
         np_key = key
+        devices = {self.device}
         if isinstance(key, tuple):
+            devices.update([subkey.device for subkey in key])
+            if len(devices) > 1:
+                raise ValueError(
+                    "Array indexing is only allowed when array to be indexed and all "
+                    "indexing arrays are on the same device."
+                )
             # Indexing self._array with array_api_strict arrays can be erroneous
             # e.g., when using non-default device
             np_key = tuple(

array_api_strict/tests/test_array_object.py

@@ -107,13 +107,14 @@ def test_indexing_arrays(device):
     device = None if device is None else Device(device)
 
     # 1D array
-    a = arange(5)
+    a = arange(5, device=device)
     idx = asarray([1, 0, 1, 2, -1], device=device)
     a_idx = a[idx]
 
     a_idx_loop = stack([a[idx[i]] for i in range(idx.shape[0])])
     assert all(a_idx == a_idx_loop)
     assert a_idx.shape == idx.shape
+    assert a.device == idx.device == a_idx.device
 
     # setitem with arrays is not allowed
     with assert_raises(IndexError):
@@ -126,20 +127,39 @@ def test_indexing_arrays(device):
     a_idx_loop = stack([a[idx[i], 1] for i in range(idx.shape[0])])
     assert all(a_idx == a_idx_loop)
     assert a_idx.shape == idx.shape
+    assert a.device == idx.device == a_idx.device
 
     # index with two arrays
     a_idx = a[idx, idx]
     a_idx_loop = stack([a[idx[i], idx[i]] for i in range(idx.shape[0])])
     assert all(a_idx == a_idx_loop)
     assert a_idx.shape == a_idx.shape
+    assert a.device == idx.device == a_idx.device
 
     # setitem with arrays is not allowed
     with assert_raises(IndexError):
         a[idx, idx] = 42
 
     # smoke test indexing with ndim > 1 arrays
     idx = idx[..., None]
-    a[idx, idx]
+    a_idx = a[idx, idx]
+    assert a.device == idx.device == a_idx.device
+
+
+def test_indexing_arrays_different_devices():
+    # Ensure indexing via array on different device errors
+    device1 = Device("CPU_DEVICE")
+    device2 = Device("device1")
+
+    a = arange(5, device=device1)
+    idx1 = asarray([1, 0, 1, 2, -1], device=device2)
+    idx2 = asarray([1, 0, 1, 2, -1], device=device1)
+
+    with pytest.raises(ValueError, match="Array indexing is only allowed when"):
+        a[idx1]
+
+    with pytest.raises(ValueError, match="Array indexing is only allowed when"):
+        a[idx1, idx2]
 
 
 def test_promoted_scalar_inherits_device():