TST: at: add test for device

Author: crusaderky

tests/test_at.py

@@ -13,6 +13,7 @@
 from array_api_extra._lib._at import _AtOp
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import array_namespace, is_writeable_array
+from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._typing import Array, SetIndex
 from array_api_extra.testing import lazy_xp_function
 
@@ -327,3 +328,13 @@ def test_gh134(xp: ModuleType, bool_mask: bool, copy: bool | None):
     idx = xp.asarray(True) if bool_mask else ()
     z = at_op(y, idx, _AtOp.SET, 1, copy=copy)
     xp_assert_equal(z, xp.asarray(1, dtype=x.dtype))
+
+
+def test_device(xp: ModuleType):
+    x = xp.asarray([1, 2, 3])
+
+    y = at(x)[:2].set(xp.asarray([4, 5]))
+    assert get_device(y) == get_device(x)
+
+    y = at(x)[xp.asarray([True, False, True])].set(4)
+    assert get_device(y) == get_device(x)