Repository: data-apis/array-api-strict

Issue:
Error with 2D array indexing using non 'CPU_DEVICE' devices
Not sure why indexing a 2D array with 'device1' or 'device2' fails. Noticed while working on https://github.com/scikit-learn/scikit-learn/pull/29431

Min reproducible example:
```python
import array_api_strict as xp

# Create device
device = xp.Device('device1')

# Create a 2D array on device1
data = xp.asarray([
    [10, 20, 30, 40],
    [50, 60, 70, 80],
    [90, 100, 110, 120]
], device=device)

# Create row and column index arrays on the same device
row_indices = xp.asarray([0, 2, 1], device=device)
col_indices = xp.asarray([1, 3, 0], device=device)

# Index the 2D array using both index arrays
result = data[row_indices, col_indices]
```

gives the following error:

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[7], line 1
----> 1 result = data[row_indices, col_indices]

File ~/miniconda3/envs/skl-array-api/lib/python3.13/site-packages/array_api_strict/_array_object.py:703, in Array.__getitem__(self, key)
    700 if isinstance(key, Array):
    701     # Indexing self._array with array_api_strict arrays can be erroneous
    702     key = key._array
--> 703 res = self._array.__getitem__(key)
    704 return self._new(res, device=self.device)

File ~/miniconda3/envs/skl-array-api/lib/python3.13/site-packages/array_api_strict/_array_object.py:167, in Array.__array__(self, dtype, copy)
    165 if _allow_array:
    166     if self._device != CPU_DEVICE:
--> 167         raise RuntimeError(f"Can not convert array on the '{self._device}' device to a Numpy array.")
    168     # copy keyword is new in 2.0.0; for older versions don't use it
    169     # retry without that keyword.
    170     if np.__version__[0] < '2':

RuntimeError: Can not convert array on the 'array_api_strict.Device('device1')' device to a Numpy array.
```

Doing the above with a 1D array works. Setting `device = xp.Device('CPU_DEVICE')` also works.

Using array-api-strict version 2.3.1

Hints from issue comments:
@ev-br: We should probably start from adding `array-api-testing` tests for device support.@lucyleeow: @ogrisel suggested the likely culprit: https://github.com/scikit-learn/scikit-learn/pull/29431#issuecomment-2743881792

Solution patch:
diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py
index 2373c89..cb2dd11 100644
--- a/array_api_strict/_array_object.py
+++ b/array_api_strict/_array_object.py
@@ -716,8 +716,24 @@ class Array:
         # Note: Only indices required by the spec are allowed. See the
         # docstring of _validate_index
         self._validate_index(key, op="getitem")
-        # Indexing self._array with array_api_strict arrays can be erroneous
-        np_key = key._array if isinstance(key, Array) else key
+        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 hasattr(subkey, "device")]
+            )
+            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(
+                subkey._array if isinstance(subkey, Array) else subkey for subkey in key
+            )
         res = self._array.__getitem__(np_key)
         return self._new(res, device=self.device)
 


Test patch:
diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py
index e24a40f..51f4f31 100644
--- a/array_api_strict/tests/test_array_object.py
+++ b/array_api_strict/tests/test_array_object.py
@@ -5,7 +5,7 @@ from numpy.testing import assert_raises, suppress_warnings
 import numpy as np
 import pytest
 
-from .. import ones, arange, reshape, asarray, result_type, all, equal
+from .. import ones, arange, reshape, asarray, result_type, all, equal, stack
 from .._array_object import Array, CPU_DEVICE, Device
 from .._dtypes import (
     _all_dtypes,
@@ -101,33 +101,40 @@ def test_validate_index():
     assert_raises(IndexError, lambda: a[idx])
 
 
-def test_indexing_arrays():
+@pytest.mark.parametrize("device", [None, "CPU_DEVICE", "device1", "device2"])
+def test_indexing_arrays(device):
     # indexing with 1D integer arrays and mixes of integers and 1D integer are allowed
+    device = None if device is None else Device(device)
 
     # 1D array
-    a = arange(5)
-    idx = asarray([1, 0, 1, 2, -1])
+    a = arange(5, device=device)
+    idx = asarray([1, 0, 1, 2, -1], device=device)
     a_idx = a[idx]
 
-    a_idx_loop = asarray([a[idx[i]] for i in range(idx.shape[0])])
+    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):
         a[idx] = 42
 
     # mixed array and integer indexing
-    a = reshape(arange(3*4), (3, 4))
-    idx = asarray([1, 0, 1, 2, -1])
+    a = reshape(arange(3*4, device=device), (3, 4))
+    idx = asarray([1, 0, 1, 2, -1], device=device)
     a_idx = a[idx, 1]
-
-    a_idx_loop = asarray([a[idx[i], 1] for i in range(idx.shape[0])])
+    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 = asarray([a[idx[i], idx[i]] for i in range(idx.shape[0])])
+    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):
@@ -135,7 +142,24 @@ def test_indexing_arrays():
 
     # 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():