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/_loss/link.py b/sklearn/_loss/link.py
index 03677c8da6139..6a42e5508fad1 100644
--- a/sklearn/_loss/link.py
+++ b/sklearn/_loss/link.py
@@ -7,11 +7,11 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from math import ulp
 
-import numpy as np
-from scipy.special import expit, logit
 from scipy.stats import gmean
 
+from sklearn.utils._array_api import _expit, _logit, get_namespace
 from sklearn.utils.extmath import softmax
 
 
@@ -41,49 +41,50 @@ def includes(self, x):
         -------
         result : bool
         """
+        xp, _ = get_namespace(x)
         if self.low_inclusive:
-            low = np.greater_equal(x, self.low)
+            low = xp.greater_equal(x, self.low)
         else:
-            low = np.greater(x, self.low)
+            low = xp.greater(x, self.low)
 
-        if not np.all(low):
+        if not xp.all(low):
             return False
 
         if self.high_inclusive:
-            high = np.less_equal(x, self.high)
+            high = xp.less_equal(x, self.high)
         else:
-            high = np.less(x, self.high)
+            high = xp.less(x, self.high)
 
         # Note: np.all returns numpy.bool_
-        return bool(np.all(high))
+        return bool(xp.all(high))
 
 
-def _inclusive_low_high(interval, dtype=np.float64):
+def _inclusive_low_high(interval):
     """Generate values low and high to be within the interval range.
 
     This is used in tests only.
 
     Returns
     -------
-    low, high : tuple
+    low, high : tuple of floats
         The returned values low and high lie within the interval.
     """
-    eps = 10 * np.finfo(dtype).eps
-    if interval.low == -np.inf:
+    eps = 10 * ulp(1)
+    if interval.low == -float("inf"):
         low = -1e10
     elif interval.low < 0:
         low = interval.low * (1 - eps) + eps
     else:
         low = interval.low * (1 + eps) + eps
 
-    if interval.high == np.inf:
+    if interval.high == float("inf"):
         high = 1e10
     elif interval.high < 0:
         high = interval.high * (1 + eps) - eps
     else:
         high = interval.high * (1 - eps) - eps
 
-    return low, high
+    return float(low), float(high)
 
 
 class BaseLink(ABC):
@@ -105,11 +106,11 @@ class BaseLink(ABC):
 
     # Usually, raw_prediction may be any real number and y_pred is an open
     # interval.
-    # interval_raw_prediction = Interval(-np.inf, np.inf, False, False)
-    interval_y_pred = Interval(-np.inf, np.inf, False, False)
+    # interval_raw_prediction = Interval(-float("inf"), float("inf"), False, False)
+    interval_y_pred = Interval(-float("inf"), float("inf"), False, False)
 
     @abstractmethod
-    def link(self, y_pred, out=None):
+    def link(self, y_pred):
         """Compute the link function g(y_pred).
 
         The link function maps (predicted) target values to raw predictions,
@@ -119,19 +120,15 @@ def link(self, y_pred, out=None):
         ----------
         y_pred : array
             Predicted target values.
-        out : array
-            A location into which the result is stored. If provided, it must
-            have a shape that the inputs broadcast to. If not provided or None,
-            a freshly-allocated array is returned.
 
         Returns
         -------
-        out : array
+        array
             Output array, element-wise link function.
         """
 
     @abstractmethod
-    def inverse(self, raw_prediction, out=None):
+    def inverse(self, raw_prediction):
         """Compute the inverse link function h(raw_prediction).
 
         The inverse link function maps raw predictions to predicted target
@@ -141,14 +138,10 @@ def inverse(self, raw_prediction, out=None):
         ----------
         raw_prediction : array
             Raw prediction values (in link space).
-        out : array
-            A location into which the result is stored. If provided, it must
-            have a shape that the inputs broadcast to. If not provided or None,
-            a freshly-allocated array is returned.
 
         Returns
         -------
-        out : array
+        array
             Output array, element-wise inverse link function.
         """
 
@@ -156,12 +149,8 @@ def inverse(self, raw_prediction, out=None):
 class IdentityLink(BaseLink):
     """The identity link function g(x)=x."""
 
-    def link(self, y_pred, out=None):
-        if out is not None:
-            np.copyto(out, y_pred)
-            return out
-        else:
-            return y_pred
+    def link(self, y_pred):
+        return y_pred  # TODO: Should we copy?
 
     inverse = link
 
@@ -169,13 +158,15 @@ def link(self, y_pred, out=None):
 class LogLink(BaseLink):
     """The log link function g(x)=log(x)."""
 
-    interval_y_pred = Interval(0, np.inf, False, False)
+    interval_y_pred = Interval(0, float("inf"), False, False)
 
-    def link(self, y_pred, out=None):
-        return np.log(y_pred, out=out)
+    def link(self, y_pred):
+        xp, _ = get_namespace(y_pred)
+        return xp.log(y_pred)
 
-    def inverse(self, raw_prediction, out=None):
-        return np.exp(raw_prediction, out=out)
+    def inverse(self, raw_prediction):
+        xp, _ = get_namespace(raw_prediction)
+        return xp.exp(raw_prediction)
 
 
 class LogitLink(BaseLink):
@@ -183,11 +174,11 @@ class LogitLink(BaseLink):
 
     interval_y_pred = Interval(0, 1, False, False)
 
-    def link(self, y_pred, out=None):
-        return logit(y_pred, out=out)
+    def link(self, y_pred):
+        return _logit(y_pred)
 
-    def inverse(self, raw_prediction, out=None):
-        return expit(raw_prediction, out=out)
+    def inverse(self, raw_prediction):
+        return _expit(raw_prediction)
 
 
 class HalfLogitLink(BaseLink):
@@ -198,13 +189,11 @@ class HalfLogitLink(BaseLink):
 
     interval_y_pred = Interval(0, 1, False, False)
 
-    def link(self, y_pred, out=None):
-        out = logit(y_pred, out=out)
-        out *= 0.5
-        return out
+    def link(self, y_pred):
+        return 0.5 * _logit(y_pred)
 
-    def inverse(self, raw_prediction, out=None):
-        return expit(2 * raw_prediction, out)
+    def inverse(self, raw_prediction):
+        return _expit(2 * raw_prediction)
 
 
 class MultinomialLogit(BaseLink):
@@ -257,20 +246,17 @@ class MultinomialLogit(BaseLink):
     interval_y_pred = Interval(0, 1, False, False)
 
     def symmetrize_raw_prediction(self, raw_prediction):
-        return raw_prediction - np.mean(raw_prediction, axis=1)[:, np.newaxis]
+        xp, _ = get_namespace(raw_prediction)
+        return raw_prediction - xp.mean(raw_prediction, axis=1)[:, None]
 
-    def link(self, y_pred, out=None):
+    def link(self, y_pred):
+        xp, _ = get_namespace(y_pred)
         # geometric mean as reference category
         gm = gmean(y_pred, axis=1)
-        return np.log(y_pred / gm[:, np.newaxis], out=out)
+        return xp.log(y_pred / gm[:, None])
 
-    def inverse(self, raw_prediction, out=None):
-        if out is None:
-            return softmax(raw_prediction, copy=True)
-        else:
-            np.copyto(out, raw_prediction)
-            softmax(out, copy=False)
-            return out
+    def inverse(self, raw_prediction):
+        return softmax(raw_prediction)
 
 
 _LINKS = {
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
index 020a758589713..0bd1869695952 100644
--- a/sklearn/utils/_array_api.py
+++ b/sklearn/utils/_array_api.py
@@ -591,18 +591,33 @@ def move_to(*arrays, xp, device):
     )
 
 
-def _expit(X, out=None, xp=None):
+def _expit(x, out=None, xp=None):
     # The out argument for exp and hence expit is only supported for numpy,
     # but not in the Array API specification.
-    xp, _ = get_namespace(X, xp=xp)
+    xp, _ = get_namespace(x, xp=xp)
     if _is_numpy_namespace(xp):
-        if out is not None:
-            special.expit(X, out=out)
-        else:
-            out = special.expit(X)
-        return out
+        return special.expit(x, out=out)
+
+    return 1.0 / (1.0 + xp.exp(-x))
 
-    return 1.0 / (1.0 + xp.exp(-X))
+
+def _logit(x, out=None, xp=None):
+    # The out argument for log and hence logit is only supported for numpy,
+    # but not in the Array API specification.
+    xp, _ = get_namespace(x, xp=xp)
+    if _is_numpy_namespace(xp):
+        return special.logit(x, out=out)
+
+    # See https://github.com/scipy/xsf/blob/e0c4d22d6ae768b39efc69586f1e8d5560a32fc5/include/xsf/log_exp.h#L30
+    def logit_v2(x):
+        s = 2 * (x - 0.5)
+        return xp.log1p(s) - xp.log1p(-s)
+
+    return xp.where(
+        xp.logical_or(x < 0.3, x > 0.65),
+        xp.log(x / (1 - x)),
+        logit_v2(x),
+    )
 
 
 def _validate_diagonal_args(array, value, xp):


Test patch:
diff --git a/sklearn/_loss/tests/test_link.py b/sklearn/_loss/tests/test_link.py
index e5a665f8d48ac..f85b223c6f1fb 100644
--- a/sklearn/_loss/tests/test_link.py
+++ b/sklearn/_loss/tests/test_link.py
@@ -1,7 +1,8 @@
 import numpy as np
 import pytest
-from numpy.testing import assert_allclose, assert_array_equal
+from numpy.testing import assert_allclose
 
+from sklearn import config_context
 from sklearn._loss.link import (
     _LINKS,
     HalfLogitLink,
@@ -9,6 +10,12 @@
     MultinomialLogit,
     _inclusive_low_high,
 )
+from sklearn.utils._array_api import (
+    _convert_to_numpy,
+    _get_namespace_device_dtype_ids,
+    yield_namespace_device_dtype_combinations,
+)
+from sklearn.utils._testing import _array_api_for_tests
 
 LINK_FUNCTIONS = list(_LINKS.values())
 
@@ -28,10 +35,10 @@ def test_interval_raises():
         Interval(0, 1, False, True),
         Interval(0, 1, True, False),
         Interval(0, 1, True, True),
-        Interval(-np.inf, np.inf, False, False),
-        Interval(-np.inf, np.inf, False, True),
-        Interval(-np.inf, np.inf, True, False),
-        Interval(-np.inf, np.inf, True, True),
+        Interval(-float("inf"), float("inf"), False, False),
+        Interval(-float("inf"), float("inf"), False, True),
+        Interval(-float("inf"), float("inf"), True, False),
+        Interval(-float("inf"), float("inf"), True, True),
         Interval(-10, -1, False, False),
         Interval(-10, -1, False, True),
         Interval(-10, -1, True, False),
@@ -39,10 +46,10 @@ def test_interval_raises():
     ],
 )
 def test_is_in_range(interval):
-    # make sure low and high are always within the interval, used for linspace
+    """Test that low and high are always within the interval used for linspace."""
     low, high = _inclusive_low_high(interval)
-
     x = np.linspace(low, high, num=10)
+
     assert interval.includes(x)
 
     # x contains lower bound
@@ -59,7 +66,7 @@ def test_is_in_range(interval):
 
 @pytest.mark.parametrize("link", LINK_FUNCTIONS)
 def test_link_inverse_identity(link, global_random_seed):
-    # Test that link of inverse gives identity.
+    """Test that link of inverse gives identity."""
     rng = np.random.RandomState(global_random_seed)
     link = link()
     n_samples, n_classes = 100, None
@@ -81,31 +88,51 @@ def test_link_inverse_identity(link, global_random_seed):
     assert_allclose(link.inverse(link.link(y_pred)), y_pred)
 
 
+@pytest.mark.parametrize(
+    "namespace, device, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
 @pytest.mark.parametrize("link", LINK_FUNCTIONS)
-def test_link_out_argument(link):
-    # Test that out argument gets assigned the result.
-    rng = np.random.RandomState(42)
+def test_link_inverse_array_api(
+    namespace, device, dtype_name, link, global_random_seed
+):
+    """Test that link and inverse link give same result for array API inputs."""
+    rng = np.random.RandomState(global_random_seed)
     link = link()
     n_samples, n_classes = 100, None
+    # The values for `raw_prediction` are limited from -20 to 20 because in the
+    # class `LogitLink` the term `expit(x)` comes very close to 1 for large
+    # positive x and therefore loses precision.
     if link.is_multiclass:
         n_classes = 10
-        raw_prediction = rng.normal(loc=0, scale=10, size=(n_samples, n_classes))
+        raw_prediction = rng.uniform(low=-20, high=20, size=(n_samples, n_classes))
         if isinstance(link, MultinomialLogit):
             raw_prediction = link.symmetrize_raw_prediction(raw_prediction)
-    else:
-        # So far, the valid interval of raw_prediction is (-inf, inf) and
-        # we do not need to distinguish.
+    elif isinstance(link, HalfLogitLink):
         raw_prediction = rng.uniform(low=-10, high=10, size=(n_samples))
+    else:
+        raw_prediction = rng.uniform(low=-20, high=20, size=(n_samples))
 
-    y_pred = link.inverse(raw_prediction, out=None)
-    out = np.empty_like(raw_prediction)
-    y_pred_2 = link.inverse(raw_prediction, out=out)
-    assert_allclose(y_pred, out)
-    assert_array_equal(out, y_pred_2)
-    assert np.shares_memory(out, y_pred_2)
-
-    out = np.empty_like(y_pred)
-    raw_prediction_2 = link.link(y_pred, out=out)
-    assert_allclose(raw_prediction, out)
-    assert_array_equal(out, raw_prediction_2)
-    assert np.shares_memory(out, raw_prediction_2)
+    xp = _array_api_for_tests(namespace, device)
+    if dtype_name != "float64":
+        raw_prediction *= 0.5  # avoid overflow
+        rtol = 1e-3 if n_classes else 1e-4
+    else:
+        rtol = 1e-8
+
+    with config_context(array_api_dispatch=True):
+        raw_prediction_xp = xp.asarray(raw_prediction.astype(dtype_name), device=device)
+        assert_allclose(
+            _convert_to_numpy(link.inverse(raw_prediction_xp), xp=xp),
+            link.inverse(raw_prediction),
+            rtol=rtol,
+        )
+
+        y_pred = link.inverse(raw_prediction)
+        y_pred_xp = xp.asarray(y_pred.astype(dtype_name), device=device)
+        assert_allclose(
+            _convert_to_numpy(link.link(y_pred_xp), xp=xp),
+            link.link(y_pred),
+            rtol=rtol,
+        )
diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py
index b300bd5aeebc6..e1abb40fc7a51 100644
--- a/sklearn/utils/tests/test_array_api.py
+++ b/sklearn/utils/tests/test_array_api.py
@@ -6,6 +6,7 @@
 import scipy
 import scipy.sparse as sp
 from numpy.testing import assert_allclose
+from scipy.special import expit, logit
 
 from sklearn._config import config_context
 from sklearn._loss import HalfMultinomialLoss
@@ -18,11 +19,13 @@
     _convert_to_numpy,
     _count_nonzero,
     _estimator_with_converted_arrays,
+    _expit,
     _fill_diagonal,
     _get_namespace_device_dtype_ids,
     _half_multinomial_loss,
     _is_numpy_namespace,
     _isin,
+    _logit,
     _logsumexp,
     _matching_numpy_dtype,
     _max_precision_float_dtype,
@@ -824,6 +827,31 @@ def test_median(namespace, device, dtype_name, axis):
     assert_allclose(result_np, _convert_to_numpy(result_xp, xp=xp))
 
 
+@pytest.mark.parametrize(
+    "namespace, device, dtype_name", yield_namespace_device_dtype_combinations()
+)
+def test_expit_logit(namespace, device, dtype_name):
+    rtol = 1e-6 if "float32" in str(dtype_name) else 1e-12
+    xp = _array_api_for_tests(namespace, device)
+
+    with config_context(array_api_dispatch=True):
+        x_np = numpy.linspace(-20, 20, 1000).astype(dtype_name)
+        x_xp = xp.asarray(x_np, device=device)
+        assert_allclose(
+            _convert_to_numpy(_expit(x_xp), xp=xp),
+            expit(x_np),
+            rtol=rtol,
+        )
+
+        x_np = numpy.linspace(0, 1, 1000).astype(dtype_name)
+        x_xp = xp.asarray(x_np, device=device)
+        assert_allclose(
+            _convert_to_numpy(_logit(x_xp), xp=xp),
+            logit(x_np),
+            rtol=rtol,
+        )
+
+
 @pytest.mark.parametrize(
     "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations()
 )