Repository: scikit-learn/scikit-learn

Issue:
# Make more of the "tools" of scikit-learn Array API compatible

🚨 🚧 This issue requires a bit of patience and experience to contribute to 🚧 🚨 

- Original issue introducing array API in scikit-learn: #22352
- array API official doc/spec: https://data-apis.org/array-api/
- scikit-learn doc: https://scikit-learn.org/dev/modules/array_api.html

Please mention this issue when you create a PR, but please don't write "closes #26024" or "fixes #26024".

scikit-learn contains lots of useful tools, in addition to the many estimators it has. For example [metrics](https://scikit-learn.org/stable/modules/classes.html#sklearn-metrics-metrics), [pipelines](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.pipeline), [pre-processing](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.preprocessing) and [mode selection](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection). These are useful to and used by people who do not necessarily use an estimator from scikit-learn. This is great.

The fact that many users install scikit-learn "just" to use `train_test_split` is a testament to how useful it is to provide easy to use tools that do the right(!) thing. Instead of everyone implementing them from scratch because it is "easy" and making mistakes along the way.

In this issue I'd like to collect and track work related to making it easier to use all these "tools" from scikit-learn even if you are not using Numpy arrays for your data. In particular thanks to the Array API standard it should be "not too much work" to make things usable with data that is in an array that conforms to the Array API standard.

There is work in #25956 and #22554 which adds the basic infrastructure needed to use "array API arrays".

The goal of this issue is to make code like the following work:
```python
>>> from sklearn.preprocessing import MinMaxScaler
>>> from sklearn import config_context
>>> from sklearn.datasets import make_classification
>>> import torch
>>> X_np, y_np = make_classification(random_state=0)
>>> X_torch = torch.asarray(X_np, device="cuda", dtype=torch.float32)
>>> y_torch = torch.asarray(y_np, device="cuda", dtype=torch.float32)

>>> with config_context(array_api_dispatch=True):
...     # For example using MinMaxScaler on PyTorch tensors
...     scale = MinMaxScaler()
...     X_trans = scale.fit_transform(X_torch, y_torch)
...     assert type(X_trans) == type(X_torch)
...     assert X_trans.device == X_torch.device
```

The first step (or maybe part of the first) is to check which of them already "just work". After that is done we can start the work (one PR per class/function) making changes.


## Guidelines for testing

General comment: most of the time when we add array API support to a function in scikit-learn, we do not touch the existing (numpy-only) tests to make sure that the PR does not change the default behavior of scikit-learn on traditional inputs when array API is not enabled.

In the case of an estimator, it can be enough to add the `array_api_support=True` estimator tag in a method named `__sklearn_tags__`. For metric functions, just register it in the `array_api_metric_checkers` in `sklearn/metrics/tests/test_common.py` to include it in the common test.

For other kinds of functions not covered by existing common tests, or when the array API support depends heavily on non-default values, it might be required to add one or more new test functions to the related module-level test file. The general testing scheme is the following:

- generate some random test data with numpy or `sklearn.datasets.make_*`;
- call the function once on the numpy inputs without enabling array API dispatch;
- convert the inputs to a namespace / device combo passed as parameter to the test;
- call the function with array API dispatching enabled (under a `with sklearn.config_context(array_api_dispatch=True)` block
- check that the results are on the same namespace and device as the input
- convert back the output to a numpy array using `_convert_to_numpy`
- compare the original / reference numpy results and the `xp` computation results converted back to numpy using `assert_allclose` or similar.

Those tests should have `array_api` somewhere in their name to makes sure that we can run all the array API compliance tests with a keyword search in the pytest command line, e.g.:

```
pytest -k array_api sklearn/some/subpackage
```

In particular, for cost reasons, our CUDA GPU CI only runs `pytest -k array_api sklearn`. So it's very important to respect this naming conventions, otherwise we will not tests all what we are supposed to test on CUDA.

More generally, look at merged array API pull requests to see how testing is typically handled.

Implementation patch:
diff --git a/sklearn/metrics/_base.py b/sklearn/metrics/_base.py
index c7668bce9fceb..9964929a446b5 100644
--- a/sklearn/metrics/_base.py
+++ b/sklearn/metrics/_base.py
@@ -10,7 +10,9 @@
 
 import numpy as np
 
+import sklearn.externals.array_api_extra as xpx
 from sklearn.utils import check_array, check_consistent_length
+from sklearn.utils._array_api import _average, _ravel, get_namespace_and_device
 from sklearn.utils.multiclass import type_of_target
 
 
@@ -19,6 +21,9 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
 
     Parameters
     ----------
+    binary_metric : callable, returns shape [n_classes]
+        The binary metric function to use.
+
     y_true : array, shape = [n_samples] or [n_samples, n_classes]
         True binary labels in binary label indicators.
 
@@ -47,9 +52,6 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
     sample_weight : array-like of shape (n_samples,), default=None
         Sample weights.
 
-    binary_metric : callable, returns shape [n_classes]
-        The binary metric function to use.
-
     Returns
     -------
     score : float or array of shape [n_classes]
@@ -57,6 +59,7 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
         classes.
 
     """
+    xp, _, _device = get_namespace_and_device(y_true, y_score, sample_weight)
     average_options = (None, "micro", "macro", "weighted", "samples")
     if average not in average_options:
         raise ValueError("average has to be one of {0}".format(average_options))
@@ -78,18 +81,23 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
 
     if average == "micro":
         if score_weight is not None:
-            score_weight = np.repeat(score_weight, y_true.shape[1])
-        y_true = y_true.ravel()
-        y_score = y_score.ravel()
+            score_weight = xp.repeat(score_weight, y_true.shape[1])
+        y_true = _ravel(y_true)
+        y_score = _ravel(y_score)
 
     elif average == "weighted":
         if score_weight is not None:
-            average_weight = np.sum(
-                np.multiply(y_true, np.reshape(score_weight, (-1, 1))), axis=0
+            #  Mixed integer and float type promotion not defined in array standard
+            y_true = xp.asarray(y_true, dtype=score_weight.dtype)
+            average_weight = xp.sum(
+                xp.multiply(y_true, xp.reshape(score_weight, (-1, 1))), axis=0
             )
         else:
-            average_weight = np.sum(y_true, axis=0)
-        if np.isclose(average_weight.sum(), 0.0):
+            average_weight = xp.sum(y_true, axis=0)
+        if xpx.isclose(
+            xp.sum(average_weight),
+            xp.asarray(0, dtype=average_weight.dtype, device=_device),
+        ):
             return 0
 
     elif average == "samples":
@@ -99,16 +107,20 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
         not_average_axis = 0
 
     if y_true.ndim == 1:
-        y_true = y_true.reshape((-1, 1))
+        y_true = xp.reshape(y_true, (-1, 1))
 
     if y_score.ndim == 1:
-        y_score = y_score.reshape((-1, 1))
+        y_score = xp.reshape(y_score, (-1, 1))
 
     n_classes = y_score.shape[not_average_axis]
-    score = np.zeros((n_classes,))
+    score = xp.zeros((n_classes,), device=_device)
     for c in range(n_classes):
-        y_true_c = y_true.take([c], axis=not_average_axis).ravel()
-        y_score_c = y_score.take([c], axis=not_average_axis).ravel()
+        y_true_c = _ravel(
+            xp.take(y_true, xp.asarray([c], device=_device), axis=not_average_axis)
+        )
+        y_score_c = _ravel(
+            xp.take(y_score, xp.asarray([c], device=_device), axis=not_average_axis)
+        )
         score[c] = binary_metric(y_true_c, y_score_c, sample_weight=score_weight)
 
     # Average the results
@@ -116,9 +128,8 @@ def _average_binary_score(binary_metric, y_true, y_score, average, sample_weight
         if average_weight is not None:
             # Scores with 0 weights are forced to be 0, preventing the average
             # score from being affected by 0-weighted NaN elements.
-            average_weight = np.asarray(average_weight)
             score[average_weight == 0] = 0
-        return float(np.average(score, weights=average_weight))
+        return float(_average(score, weights=average_weight, xp=xp))
     else:
         return score
 
diff --git a/sklearn/metrics/_ranking.py b/sklearn/metrics/_ranking.py
index 782f5c0fc7dbe..8712c63f0780a 100644
--- a/sklearn/metrics/_ranking.py
+++ b/sklearn/metrics/_ranking.py
@@ -31,6 +31,7 @@
 from sklearn.utils._array_api import (
     _max_precision_float_dtype,
     get_namespace_and_device,
+    move_to,
     size,
 )
 from sklearn.utils._encode import _encode, _unique
@@ -225,25 +226,36 @@ def average_precision_score(
     >>> average_precision_score(y_true, y_scores)
     0.77
     """
+    xp, _, device = get_namespace_and_device(y_score)
+    y_true, sample_weight = move_to(y_true, sample_weight, xp=xp, device=device)
+
+    if sample_weight is not None:
+        sample_weight = column_or_1d(sample_weight)
 
     def _binary_uninterpolated_average_precision(
-        y_true, y_score, pos_label=1, sample_weight=None
+        y_true,
+        y_score,
+        pos_label=1,
+        sample_weight=None,
+        xp=xp,
     ):
         precision, recall, _ = precision_recall_curve(
-            y_true, y_score, pos_label=pos_label, sample_weight=sample_weight
+            y_true,
+            y_score,
+            pos_label=pos_label,
+            sample_weight=sample_weight,
         )
         # Return the step function integral
         # The following works because the last entry of precision is
         # guaranteed to be 1, as returned by precision_recall_curve.
         # Due to numerical error, we can get `-0.0` and we therefore clip it.
-        return float(max(0.0, -np.sum(np.diff(recall) * np.array(precision)[:-1])))
+        return float(max(0.0, -xp.sum(xp.diff(recall) * precision[:-1])))
 
     y_type = type_of_target(y_true, input_name="y_true")
-
-    present_labels = np.unique(y_true)
+    present_labels = xp.unique_values(y_true)
 
     if y_type == "binary":
-        if len(present_labels) == 2 and pos_label not in present_labels:
+        if present_labels.shape[0] == 2 and pos_label not in present_labels:
             raise ValueError(
                 f"pos_label={pos_label} is not a valid label. It should be "
                 f"one of {present_labels}"
@@ -270,7 +282,7 @@ def _binary_uninterpolated_average_precision(
             )
 
     average_precision = partial(
-        _binary_uninterpolated_average_precision, pos_label=pos_label
+        _binary_uninterpolated_average_precision, pos_label=pos_label, xp=xp
     )
     return _average_binary_score(
         average_precision, y_true, y_score, average, sample_weight=sample_weight
@@ -686,6 +698,8 @@ class scores must correspond to the order of ``labels``,
     y_type = type_of_target(y_true, input_name="y_true")
     y_true = check_array(y_true, ensure_2d=False, dtype=None)
     y_score = check_array(y_score, ensure_2d=False)
+    if sample_weight is not None:
+        sample_weight = column_or_1d(sample_weight)
 
     if y_type == "multiclass" or (
         y_type == "binary" and y_score.ndim == 2 and y_score.shape[1] > 2
@@ -1142,7 +1156,7 @@ def precision_recall_curve(
             "No positive class found in y_true, "
             "recall is set to one for all thresholds."
         )
-        recall = xp.full(tps.shape, 1.0)
+        recall = xp.full(tps.shape, 1.0, device=device)
     else:
         recall = tps / tps[-1]
 


Test patch:
diff --git a/sklearn/metrics/tests/test_common.py b/sklearn/metrics/tests/test_common.py
index 353dfebf93bcf..0b7d8b474cec1 100644
--- a/sklearn/metrics/tests/test_common.py
+++ b/sklearn/metrics/tests/test_common.py
@@ -2098,6 +2098,18 @@ def check_array_api_multiclass_classification_metric(
     y_true_np = np.array([0, 1, 2, 3])
     y_pred_np = np.array([0, 1, 0, 2])
 
+    if metric.__name__ == "average_precision_score":
+        # we need y_pred_nd to be of shape (n_samples, n_classes)
+        y_pred_np = np.array(
+            [
+                [0.7, 0.2, 0.05, 0.05],
+                [0.1, 0.8, 0.05, 0.05],
+                [0.1, 0.1, 0.7, 0.1],
+                [0.05, 0.05, 0.1, 0.8],
+            ],
+            dtype=dtype_name,
+        )
+
     additional_params = {
         "average": ("micro", "macro", "weighted"),
         "beta": (0.2, 0.5, 0.8),
@@ -2299,6 +2311,11 @@ def check_array_api_metric_pairwise(metric, array_namespace, device, dtype_name)
         check_array_api_multiclass_classification_metric,
         check_array_api_multilabel_classification_metric,
     ],
+    average_precision_score: [
+        check_array_api_binary_classification_metric,
+        check_array_api_multiclass_classification_metric,
+        check_array_api_multilabel_classification_metric,
+    ],
     balanced_accuracy_score: [
         check_array_api_binary_classification_metric,
         check_array_api_multiclass_classification_metric,