Make common test_check_targets_array_api (scikit-learn#34487)

Author: lucyleeow

sklearn/metrics/_classification.py

@@ -124,8 +124,10 @@ def _check_targets(y_true, y_pred, sample_weight=None, xp=None, device=None):
     """
     if xp is None or device is None:
         xp, _, device = get_namespace_and_device(y_pred, xp=xp)
-    # Ensure `y_true` and `sample_weight` are on the `y_pred` device and namespace.
-    y_true, sample_weight = move_to(y_true, sample_weight, xp=xp, device=device)
+    # Move to specified namespace and device, or the namespace and device of `y_pred`
+    y_true, y_pred, sample_weight = move_to(
+        y_true, y_pred, sample_weight, xp=xp, device=device
+    )
     y_true, y_pred = attach_unique(y_true, y_pred)
 
     check_consistent_length(y_true, y_pred, sample_weight)

sklearn/metrics/tests/test_common.py

@@ -62,6 +62,8 @@
     zero_one_loss,
 )
 from sklearn.metrics._base import _average_binary_score
+from sklearn.metrics._classification import _check_targets
+from sklearn.metrics._regression import _check_reg_targets
 from sklearn.metrics.pairwise import (
     additive_chi2_kernel,
     chi2_kernel,
@@ -2049,6 +2051,44 @@ def test_metrics_pos_label_error_str(metric, y_pred_threshold, dtype_y_str):
         metric(y1, y2)
 
 
+@pytest.mark.parametrize("check", (_check_targets, _check_reg_targets))
+def test_check_targets_array_api(check):
+    """Ensure `check` returns arrays in the correct namespace and device."""
+    xp, _ = _array_api_for_tests("array_api_strict")
+    y_pred = np.asarray([[0, 0, 0, 1], [1, 0, 1, 1], [0, 0, 0, 1]])
+    y_pred_xp = xp.asarray(y_pred)
+
+    kwarg = (
+        {"multioutput": np.asarray([0.2, 0.2, 0.25, 0.35])}
+        if check is _check_reg_targets
+        else {}
+    )
+    with config_context(array_api_dispatch=True):
+        # Do not pass `xp` and `device`
+        outputs = check(
+            y_true=np.asarray([[1, 0, 0, 1], [0, 1, 1, 1], [1, 1, 0, 1]]),
+            y_pred=y_pred_xp,
+            sample_weight=np.asarray([1, 1, 2]),
+            **kwarg,
+        )
+        for output in outputs[1:]:
+            assert get_namespace(output)[0] == xp
+            assert array_api_device(output) == array_api_device(y_pred_xp)
+
+        # Pass `xp` and `device`
+        outputs = check(
+            y_true=np.asarray([[1, 0, 0, 1], [0, 1, 1, 1], [1, 1, 0, 1]]),
+            y_pred=y_pred,
+            sample_weight=np.asarray([1, 1, 2]),
+            xp=xp,
+            device=array_api_device(y_pred_xp),
+            **kwarg,
+        )
+        for output in outputs[1:]:
+            assert get_namespace(output)[0] == xp
+            assert array_api_device(output) == array_api_device(y_pred_xp)
+
+
 def check_array_api_metric(
     metric, array_namespace, device_name, dtype_name, a_np, b_np, **metric_kwargs
 ):

sklearn/metrics/tests/test_regression.py

@@ -6,7 +6,6 @@
 from scipy import optimize
 from scipy.special import factorial, xlogy
 
-from sklearn._config import config_context
 from sklearn.dummy import DummyRegressor
 from sklearn.exceptions import UndefinedMetricWarning
 from sklearn.metrics import (
@@ -29,9 +28,7 @@
 )
 from sklearn.metrics._regression import _check_reg_targets
 from sklearn.model_selection import GridSearchCV
-from sklearn.utils._array_api import device, get_namespace
 from sklearn.utils._testing import (
-    _array_api_for_tests,
     assert_almost_equal,
     assert_array_almost_equal,
     assert_array_equal,
@@ -367,38 +364,6 @@ def test__check_reg_targets_single_output_error():
         _check_reg_targets([1, 2, 3, 4], [1.5, 2.5, 3.5, 4.5], None, multioutput=[0.5])
 
 
-def test__check_reg_targets_array_api():
-    """Check arrays are returned in the correct namespace."""
-    xp, _ = _array_api_for_tests("array_api_strict")
-    y_pred = np.asarray([[0, 0, 0, 1], [1, 0, 1, 1], [0, 0, 0, 1]])
-    y_pred_xp = xp.asarray(y_pred)
-
-    with config_context(array_api_dispatch=True):
-        # Do not pass `xp` and `device`
-        outputs = _check_reg_targets(
-            y_true=np.asarray([[1, 0, 0, 1], [0, 1, 1, 1], [1, 1, 0, 1]]),
-            y_pred=y_pred_xp,
-            sample_weight=np.asarray([1, 1, 2]),
-            multioutput=np.asarray([0.2, 0.2, 0.25, 0.35]),
-        )
-        for output in outputs[1:]:
-            assert get_namespace(output)[0] == xp
-            assert device(output) == device(y_pred_xp)
-
-        # Pass `xp` and `device`
-        outputs = _check_reg_targets(
-            y_true=np.asarray([[1, 0, 0, 1], [0, 1, 1, 1], [1, 1, 0, 1]]),
-            y_pred=y_pred,
-            sample_weight=np.asarray([1, 1, 2]),
-            multioutput=np.asarray([0.2, 0.2, 0.25, 0.35]),
-            xp=xp,
-            device=device(y_pred_xp),
-        )
-        for output in outputs[1:]:
-            assert get_namespace(output)[0] == xp
-            assert device(output) == device(y_pred_xp)
-
-
 def test_regression_multioutput_array():
     y_true = [[1, 2], [2.5, -1], [4.5, 3], [5, 7]]
     y_pred = [[1, 1], [2, -1], [5, 4], [5, 6.5]]