Merge pull request data-apis#209 from ev-br/setitem_device

ENH: make __setitem__ raise on a device mismatch

Author: ev-br

array_api_strict/_array_object.py

@@ -962,6 +962,10 @@ def __setitem__(
         other = value
         if isinstance(value, (bool, int, float, complex)):
             other = self._promote_scalar(value)
+        else:
+            if value.device != self.device:
+                raise ValueError(f"mismatched devices: {self.device = } != {value.device =}.")
+
         dt = _result_type(self.dtype, other.dtype)
         if dt != self.dtype:
             raise TypeError(f"mismatched dtypes: {self.dtype = } and {other.dtype = }")

array_api_strict/tests/test_array_object.py

@@ -221,6 +221,14 @@ def test_setitem_invalid_promotions():
         a[0] = asarray(3.5j, dtype=complex128)
 
 
+def test_setitem_device_transfer():
+    a = arange(3)
+    b = arange(4, 1, -1, device=Device('device1'))
+
+    with pytest.raises(ValueError):
+        a[:] = b[:]
+
+
 def test_promoted_scalar_inherits_device():
     device1 = Device("device1")
     x = asarray([1., 2, 3], device=device1)