FEA array API support for LogisticRegressionCV (scikit-learn#33906)

Author: OmarManzoor

doc/modules/array_api.rst

@@ -163,6 +163,7 @@ Estimators
   `svd_solver="randomized"` (`svd_solver="randomized"` only if `power_iteration_normalizer="QR"`))
 - :class:`kernel_approximation.Nystroem`
 - :class:`linear_model.LogisticRegression` (with `solver="lbfgs"` and `solver="newton-cg"`)
+- :class:`linear_model.LogisticRegressionCV` (with `solver="lbfgs"` and `solver="newton-cg"`)
 - :class:`linear_model.PoissonRegressor` (with `solver="lbfgs"`)
 - :class:`linear_model.Ridge` (with `solver="svd"`)
 - :class:`linear_model.RidgeCV` (see :ref:`device_support_for_float64`)

doc/whats_new/upcoming_changes/array-api/33906.feature.rst

@@ -0,0 +1,4 @@
+- :class:`linear_model.LogisticRegressionCV` now supports array API compatible inputs
+  with `solver="lbfgs"` and `solver="newton-cg"` and when the underlying scoring
+  function is also compatible with the array API.
+  By :user:`Omar Salman <OmarManzoor>` and :user:`Christian Lorentzen <lorentzenchr>`.

sklearn/linear_model/_logistic.py

@@ -2879,12 +2879,10 @@ def test_logistic_regression_array_api_compliance(
     # the intercept would be non-zero.
     lr_params = dict(C=1e-2, solver=solver, max_iter=500, class_weight=class_weight)
     lr_params["tol"] = 1e-6 if dtype_name == "float32" else 1e-12
-    with warnings.catch_warnings():
-        # Make sure that we converge in the reference fit.
-        lr_np = LogisticRegression(**lr_params).fit(
-            X_np, y_np, sample_weight=sample_weight
-        )
-        assert lr_np.n_iter_ < lr_np.max_iter
+
+    lr_np = LogisticRegression(**lr_params).fit(X_np, y_np, sample_weight=sample_weight)
+    # Make sure that the reference fit converged.
+    assert lr_np.n_iter_ < lr_np.max_iter
 
     # Test that C was not too large for meaningful testing.
     assert np.abs(lr_np.coef_).max() > 0.1
@@ -2900,14 +2898,11 @@ def test_logistic_regression_array_api_compliance(
         rtol = 1e-4 if dtype_name == "float32" else 1e-8
 
     with config_context(array_api_dispatch=True):
-        with warnings.catch_warnings():
-            # Make sure that we converge when using the namespace/device
-            # specific fit.
-            warnings.simplefilter("error", ConvergenceWarning)
-            lr_xp = LogisticRegression(**lr_params).fit(
-                X_xp, y_xp_or_np, sample_weight=sample_weight
-            )
-
+        lr_xp = LogisticRegression(**lr_params).fit(
+            X_xp, y_xp_or_np, sample_weight=sample_weight
+        )
+        # Make sure that we converge when using the namespace/device
+        # specific fit.
         assert lr_xp.n_iter_.shape == lr_np.n_iter_.shape
         assert int(lr_xp.n_iter_[0]) < lr_xp.max_iter
 
@@ -2946,6 +2941,128 @@ def test_logistic_regression_array_api_compliance(
         assert_array_equal(prediction_xp, prediction_np)
 
 
+@pytest.mark.parametrize("solver", ["lbfgs", "newton-cg"])
+@pytest.mark.parametrize("refit", [True, False])
+@pytest.mark.parametrize("binary", [False, True])
+@pytest.mark.parametrize("use_str_y", [False, True])
+@pytest.mark.parametrize("use_sample_weight", [False, True])
+@pytest.mark.parametrize(
+    "array_namespace, device_name, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+)
+@pytest.mark.filterwarnings("error::sklearn.exceptions.ConvergenceWarning")
+def test_logistic_regression_cv_array_api_compliance(
+    solver,
+    refit,
+    binary,
+    use_str_y,
+    use_sample_weight,
+    array_namespace,
+    device_name,
+    dtype_name,
+):
+    xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
+    X_np = iris.data.astype(dtype_name, copy=True)
+    n_samples, _ = X_np.shape
+    X_xp = xp.asarray(X_np, device=device)
+    if use_str_y:
+        if binary:
+            target = (iris.target > 0).astype(np.int64)
+            target = np.array(["setosa", "not-setosa"])[target]
+            class_weight = {"setosa": 1.0, "not-setosa": 3.0}
+        else:
+            target = iris.target_names[iris.target]
+            class_weight = {"virginica": 1.0, "setosa": 2.0, "versicolor": 3.0}
+        y_np = target.copy()
+        y_xp_or_np = np.asarray(y_np, copy=True)
+    else:
+        if binary:
+            target = (iris.target > 0).astype(np.int64)
+            class_weight = {0: 1.0, 1: 3.0}
+        else:
+            target = iris.target
+            class_weight = {0: 1.0, 1: 2.0, 2: 3.0}
+        y_np = target.astype(dtype_name)
+        y_xp_or_np = xp.asarray(y_np, device=device)
+
+    if use_sample_weight:
+        sample_weight = (
+            np.random.default_rng(0)
+            .uniform(-1, 5, size=n_samples)
+            .clip(0, None)
+            .astype(dtype_name)
+        )
+    else:
+        sample_weight = None
+
+    cv = StratifiedKFold(2, shuffle=False)
+    precomputed_folds = list(cv.split(X_np, y_np))
+    atol = _atol_for_type(dtype_name) * 10
+    if dtype_name == "float32":
+        # TODO: this test is currently very sensitive to the value
+        # of the solver tol, particularly for "lbgfs". The rtol value
+        # is also quite high. Therefore try to investigate the numerical
+        # stability of the convergence with "lbgfs" in a dedicated
+        # follow up issue and PR.
+        if solver == "lbfgs":
+            solver_tol = 5e-3
+            rtol = 6e-2
+        else:
+            solver_tol = 1e-5
+            rtol = 5e-3
+    else:
+        solver_tol = 1e-10
+        rtol = 5e-5
+
+    lr_cv_params = dict(
+        Cs=[0.01, 0.001],
+        cv=precomputed_folds,
+        solver=solver,
+        tol=solver_tol,
+        max_iter=200,
+        class_weight=class_weight,
+        scoring="neg_log_loss",
+        use_legacy_attributes=False,
+        refit=refit,
+    )
+
+    lr_cv_np = LogisticRegressionCV(**lr_cv_params).fit(
+        X_np, y_np, sample_weight=sample_weight
+    )
+    # Make sure that the reference fit converged.
+    assert np.max(lr_cv_np.n_iter_) < lr_cv_np.max_iter
+
+    # Test that C was not too large for meaningful testing.
+    assert np.abs(lr_cv_np.coef_).max() > 0.1
+
+    prediction_np = lr_cv_np.predict(X_np)
+    score_np = lr_cv_np.score(X_np, y_np)
+    xp, _ = _array_api_for_tests(array_namespace, device_name)
+    with config_context(array_api_dispatch=True):
+        lr_cv_xp = LogisticRegressionCV(**lr_cv_params).fit(
+            X_xp, y_xp_or_np, sample_weight=sample_weight
+        )
+        assert lr_cv_xp.n_iter_.shape == lr_cv_np.n_iter_.shape
+        assert xp.max(lr_cv_xp.n_iter_) < lr_cv_xp.max_iter
+
+        for attr_name in ("scores_", "coefs_paths_", "coef_", "intercept_"):
+            attr_xp = getattr(lr_cv_xp, attr_name)
+            attr_np = getattr(lr_cv_np, attr_name)
+            assert_allclose(
+                move_to(attr_xp, xp=np, device="cpu"), attr_np, rtol=rtol, atol=atol
+            )
+            assert attr_xp.dtype == X_xp.dtype
+            assert array_api_device(attr_xp) == array_api_device(X_xp)
+
+        prediction_xp = lr_cv_xp.predict(X_xp)
+        if not use_str_y:
+            prediction_xp = move_to(prediction_xp, xp=np, device="cpu")
+        assert_array_equal(prediction_xp, prediction_np)
+
+        score_xp = lr_cv_xp.score(X_xp, y_xp_or_np)
+        assert_allclose(score_xp, score_np, rtol=rtol, atol=atol)
+
+
 # TODO(1.10): remove when penalty is removed
 @pytest.mark.filterwarnings("ignore:'penalty' was deprecated")
 @pytest.mark.parametrize("penalty, l1_ratio", [("l1", 0.0), ("l2", 1.0)])

sklearn/linear_model/tests/test_logistic.py

@@ -1357,3 +1357,31 @@ def _matching_numpy_dtype(X, xp=None):
     reversed_dtypes_dict = {dtype: name for name, dtype in dtypes_dict.items()}
     dtype_name = reversed_dtypes_dict[X.dtype]
     return np_compat.__array_namespace_info__().dtypes()[dtype_name]
+
+
+def _swapaxes(array, axis1, axis2, /, xp=None):
+    xp, _ = get_namespace(array)
+    if hasattr(xp, "swapaxes"):
+        return xp.swapaxes(array, axis1, axis2)
+
+    ndim = array.ndim
+    a1 = axis1 if axis1 >= 0 else axis1 + ndim
+    a2 = axis2 if axis2 >= 0 else axis2 + ndim
+    axes = list(range(ndim))
+    axes[a1], axes[a2] = axes[a2], axes[a1]
+    return xp.permute_dims(array, axes=tuple(axes))
+
+
+def _unravel_index(indices, shape, /, xp=None):
+    # TODO: remove this and use the respective array-api-extra function when
+    # version 0.11.1 is released.
+    # https://data-apis.org/array-api-extra/generated/array_api_extra.unravel_index.html
+    xp, _ = get_namespace(indices)
+    if hasattr(xp, "unravel_index"):
+        return xp.unravel_index(indices, shape)
+
+    coordinates = []
+    for dim in reversed(shape):
+        coordinates.append(indices % dim)
+        indices = indices // dim
+    return tuple(reversed(coordinates))

sklearn/utils/_array_api.py

@@ -9,6 +9,7 @@
 import numpy as np
 
 from sklearn.base import is_classifier
+from sklearn.utils._array_api import _ravel, get_namespace
 from sklearn.utils.multiclass import type_of_target
 from sklearn.utils.validation import _check_response_method, check_is_fitted
 
@@ -57,7 +58,8 @@ def _process_predict_proba(*, y_pred, target_type, classes, pos_label):
         )
 
     if target_type == "binary":
-        col_idx = np.flatnonzero(classes == pos_label)[0]
+        xp, _ = get_namespace(classes)
+        col_idx = int(xp.nonzero(_ravel(classes == pos_label))[0][0])
         return y_pred[:, col_idx]
     elif target_type == "multilabel-indicator":
         # Use a compress format of shape `(n_samples, n_output)`.

sklearn/utils/_response.py

@@ -1164,6 +1164,12 @@ def _check_array_api_core(
     # Fitted attributes which are arrays must have the same namespace as `X`,
     # except `classes_`, to allow it to be string when `y` is string.
     for attribute_name, attribute_value in array_attributes.items():
+        # TODO(1.10): remove with deprecation of default l1_ratios
+        if attribute_name == "l1_ratios_":
+            # l1_ratios_ can be None and so it is skipped in this test
+            # because it is not associated with any array namespace or device.
+            continue
+
         est_xp_attr = getattr(est_xp, attribute_name)
         # `classes_` should be in same ns and device as `y`
         expected_xp, expected_ns = (

sklearn/utils/estimator_checks.py

@@ -30,6 +30,7 @@
     _nanmean,
     _nanmin,
     _ravel,
+    _swapaxes,
     _validate_diagonal_args,
     check_same_namespace,
     get_namespace,
@@ -1042,3 +1043,17 @@ def test_matching_numpy_dtype(namespace, device_name, dtype_name):
     with config_context(array_api_dispatch=True):
         ret_dtype = _matching_numpy_dtype(X_xp, xp=xp)
     assert ret_dtype == X_np.dtype
+
+
+@pytest.mark.parametrize(
+    "namespace, device_name, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+)
+def test_swapaxes(namespace, device_name, dtype_name):
+    xp, device = _array_api_for_tests(namespace, device_name, dtype_name)
+    X_np = numpy.arange(10).reshape(5, 2).astype(dtype_name)
+    X_xp = xp.asarray(X_np, device=device)
+    result_np = numpy.swapaxes(X_np, 0, 1)
+    with config_context(array_api_dispatch=True):
+        result_xp = _swapaxes(X_xp, 0, 1)
+    assert_array_equal(move_to(result_xp, xp=numpy, device="cpu"), result_np)