TST fix check_array_api_input device check (scikit-learn#31814)

Co-authored-by: Loïc Estève <loic.esteve@ymail.com>

Author: StefanieSenger

sklearn/model_selection/tests/test_split.py

@@ -1357,11 +1357,11 @@ def test_array_api_train_test_split(
         assert get_namespace(y_train_xp)[0] == get_namespace(y_xp)[0]
         assert get_namespace(y_test_xp)[0] == get_namespace(y_xp)[0]
 
-    # Check device and dtype is preserved on output
-    assert array_api_device(X_train_xp) == array_api_device(X_xp)
-    assert array_api_device(y_train_xp) == array_api_device(y_xp)
-    assert array_api_device(X_test_xp) == array_api_device(X_xp)
-    assert array_api_device(y_test_xp) == array_api_device(y_xp)
+        # Check device and dtype is preserved on output
+        assert array_api_device(X_train_xp) == array_api_device(X_xp)
+        assert array_api_device(y_train_xp) == array_api_device(y_xp)
+        assert array_api_device(X_test_xp) == array_api_device(X_xp)
+        assert array_api_device(y_test_xp) == array_api_device(y_xp)
 
     assert X_train_xp.dtype == X_xp.dtype
     assert y_train_xp.dtype == y_xp.dtype

sklearn/utils/estimator_checks.py

@@ -1093,7 +1093,8 @@ def check_array_api_input(
             f"got {attribute_ns}"
         )
 
-        assert array_device(est_xp_param) == array_device(X_xp)
+        with config_context(array_api_dispatch=True):
+            assert array_device(est_xp_param) == array_device(X_xp)
 
         est_xp_param_np = _convert_to_numpy(est_xp_param, xp=xp)
         if check_values:
@@ -1180,7 +1181,9 @@ def check_array_api_input(
             f"got {result_ns}."
         )
 
-        assert array_device(result_xp) == array_device(X_xp)
+        with config_context(array_api_dispatch=True):
+            assert array_device(result_xp) == array_device(X_xp)
+
         result_xp_np = _convert_to_numpy(result_xp, xp=xp)
 
         if check_values:
@@ -1205,7 +1208,8 @@ def check_array_api_input(
                 f" {input_ns}, got {inverse_result_ns}."
             )
 
-            assert array_device(invese_result_xp) == array_device(X_xp)
+            with config_context(array_api_dispatch=True):
+                assert array_device(invese_result_xp) == array_device(X_xp)
 
             invese_result_xp_np = _convert_to_numpy(invese_result_xp, xp=xp)
             if check_values:

sklearn/utils/tests/test_array_api.py

@@ -166,10 +166,10 @@ def test_average(
     with config_context(array_api_dispatch=True):
         result = _average(array_in, axis=axis, weights=weights, normalize=normalize)
 
-    if np_version < parse_version("2.0.0") or np_version >= parse_version("2.1.0"):
-        # NumPy 2.0 has a problem with the device attribute of scalar arrays:
-        # https://github.com/numpy/numpy/issues/26850
-        assert device(array_in) == device(result)
+        if np_version < parse_version("2.0.0") or np_version >= parse_version("2.1.0"):
+            # NumPy 2.0 has a problem with the device attribute of scalar arrays:
+            # https://github.com/numpy/numpy/issues/26850
+            assert device(array_in) == device(result)
 
     result = _convert_to_numpy(result, xp)
     assert_allclose(result, expected, atol=_atol_for_type(dtype_name))