Fix _safe_indexing with non integer arrays on array API inputs (#32840)

ogrisel · lesteve · commit 764249a255ba · 2025-12-09T17:12:56.000+01:00
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -469,7 +469,7 @@ def move_to(*arrays, xp, device):
     `array` may contain `None` entries, these are left unchanged.
 
     Sparse arrays are accepted (as pass through) if the reference namespace is
-    Numpy, in which case they are returned unchanged. Otherwise a `TypeError`
+    NumPy, in which case they are returned unchanged. Otherwise a `TypeError`
     is raised.
 
     Parameters
diff --git a/sklearn/utils/_indexing.py b/sklearn/utils/_indexing.py
@@ -36,8 +36,24 @@ def _array_indexing(array, key, key_dtype, axis):
     """Index an array or scipy.sparse consistently across NumPy version."""
     xp, is_array_api, device_ = get_namespace_and_device(array)
     if is_array_api:
-        key = move_to(key, xp=xp, device=device_)
-        return xp.take(array, key, axis=axis)
+        if hasattr(key, "shape"):
+            key = move_to(key, xp=xp, device=device_)
+        elif isinstance(key, (int, slice)):
+            # Passthrough for valid __getitem__ inputs as noted in the array
+            # API spec.
+            pass
+        else:
+            key = xp.asarray(key, device=device_)
+
+        if hasattr(key, "dtype"):
+            if xp.isdtype(key.dtype, "integral"):
+                return xp.take(array, key, axis=axis)
+            elif xp.isdtype(key.dtype, "bool"):
+                # Array API does not support boolean indexing for n-dim arrays
+                # yet hence the need to turn to equivalent integer indexing.
+                indices = xp.arange(array.shape[axis], device=device_)
+                return xp.take(array, indices[key], axis=axis)
+
     if issparse(array) and key_dtype == "bool":
         key = np.asarray(key)
     if isinstance(key, tuple):
diff --git a/sklearn/utils/tests/test_indexing.py b/sklearn/utils/tests/test_indexing.py
@@ -10,7 +10,10 @@
 from sklearn.externals._packaging.version import parse as parse_version
 from sklearn.utils import _safe_indexing, resample, shuffle
 from sklearn.utils._array_api import (
+    _convert_to_numpy,
     _get_namespace_device_dtype_ids,
+    device,
+    move_to,
     yield_namespace_device_dtype_combinations,
 )
 from sklearn.utils._indexing import (
@@ -22,6 +25,7 @@
 from sklearn.utils._testing import (
     _array_api_for_tests,
     _convert_container,
+    assert_allclose,
     assert_allclose_dense_sparse,
     assert_array_equal,
     skip_if_array_api_compat_not_configured,
@@ -108,22 +112,22 @@ def test_determine_key_type_slice_error():
 
 @skip_if_array_api_compat_not_configured
 @pytest.mark.parametrize(
-    "array_namespace, device, dtype_name",
+    "array_namespace, device_, dtype_name",
     yield_namespace_device_dtype_combinations(),
     ids=_get_namespace_device_dtype_ids,
 )
-def test_determine_key_type_array_api(array_namespace, device, dtype_name):
-    xp = _array_api_for_tests(array_namespace, device)
+def test_determine_key_type_array_api(array_namespace, device_, dtype_name):
+    xp = _array_api_for_tests(array_namespace, device_)
 
     with sklearn.config_context(array_api_dispatch=True):
-        int_array_key = xp.asarray([1, 2, 3])
+        int_array_key = xp.asarray([1, 2, 3], device=device_)
         assert _determine_key_type(int_array_key) == "int"
 
-        bool_array_key = xp.asarray([True, False, True])
+        bool_array_key = xp.asarray([True, False, True], device=device_)
         assert _determine_key_type(bool_array_key) == "bool"
 
         try:
-            complex_array_key = xp.asarray([1 + 1j, 2 + 2j, 3 + 3j])
+            complex_array_key = xp.asarray([1 + 1j, 2 + 2j, 3 + 3j], device=device_)
         except TypeError:
             # Complex numbers are not supported by all Array API libraries.
             complex_array_key = None
@@ -133,6 +137,42 @@ def test_determine_key_type_array_api(array_namespace, device, dtype_name):
                 _determine_key_type(complex_array_key)
 
 
+@skip_if_array_api_compat_not_configured
+@pytest.mark.parametrize(
+    "array_namespace, device_, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
+@pytest.mark.parametrize(
+    "indexing_key",
+    (
+        0,
+        -1,
+        [1, 3],
+        np.array([1, 3]),
+        slice(1, 2),
+        [True, False, True, True],
+        np.asarray([False, False, False, False]),
+    ),
+)
+@pytest.mark.parametrize("axis", [0, 1])
+def test_safe_indexing_array_api_support(
+    array_namespace, device_, dtype_name, indexing_key, axis
+):
+    xp = _array_api_for_tests(array_namespace, device_)
+
+    array_to_index_np = np.arange(16).reshape(4, 4)
+    expected_result = _safe_indexing(array_to_index_np, indexing_key, axis=axis)
+    array_to_index_xp = move_to(array_to_index_np, xp=xp, device=device_)
+
+    with sklearn.config_context(array_api_dispatch=True):
+        indexed_array_xp = _safe_indexing(array_to_index_xp, indexing_key, axis=axis)
+        assert device(indexed_array_xp) == device(array_to_index_xp)
+        assert indexed_array_xp.dtype == array_to_index_xp.dtype
+
+    assert_allclose(_convert_to_numpy(indexed_array_xp, xp=xp), expected_result)
+
+
 @pytest.mark.parametrize(
     "array_type", ["list", "array", "sparse", "dataframe", "polars", "pyarrow"]
 )
