ENH add Array API to newton-cg in LogisticRegression (scikit-learn#34412)

Co-authored-by: Omar Salman <omar.salman@arbisoft.com>

Author: lorentzenchr

doc/modules/array_api.rst

@@ -162,7 +162,7 @@ Estimators
 - :class:`decomposition.PCA` (with `svd_solver="full"`, `svd_solver="covariance_eigh"`, or
   `svd_solver="randomized"` (`svd_solver="randomized"` only if `power_iteration_normalizer="QR"`))
 - :class:`kernel_approximation.Nystroem`
-- :class:`linear_model.LogisticRegression` (with `solver="lbfgs"`)
+- :class:`linear_model.LogisticRegression` (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/34412.enhancement.rst

@@ -0,0 +1,3 @@
+ :class:`linear_model.LogisticRegression` now supports array API compatible inputs
+  with `solver="newton-cg"`.
+  By :user:`Christian Lorentzen <lorentzenchr>` and :user:`Omar Salman <OmarManzoor>`.

doc/whats_new/upcoming_changes/sklearn.linear_model/34412.fix.rst

@@ -0,0 +1,6 @@
+- The `solver="newton-cg"` for :class:`linear_model.LogisticRegression`,
+  :class:`linear_model.PoissonRegressor` and other GLMs is now more robust. Many
+  convergence issues have been fixed, e.g. for certain ill-conditioned problems with
+  low regularization and nearly collinear features. The fix is about treating small
+  or even negative curvature.
+  By :user:`Christian Lorentzen <lorentzenchr>`.

sklearn/_loss/tests/test_loss.py

@@ -1411,6 +1411,9 @@ def test_loss_array_api(
     device_name,
     dtype_name,
 ):
+    if loss_class == HalfMultinomialLoss and method_name == "gradient_hessian":
+        pytest.skip("Not implemented")
+
     def _assert_array_api_result(
         result_xp, result_np, raw_prediction_xp, xp, rtol, atol
     ):

sklearn/linear_model/_linear_loss.py

@@ -9,6 +9,8 @@
 from scipy import sparse
 
 from sklearn.utils._array_api import (
+    _is_numpy_namespace,
+    _ravel,
     get_namespace,
     get_namespace_and_device,
     move_to,
@@ -172,7 +174,11 @@ def weight_intercept(self, coef):
         else:
             # reshape to (n_classes, n_dof)
             if coef.ndim == 1:
-                weights = coef.reshape((self.base_loss.n_classes, -1), order="F")
+                xp, _ = get_namespace(coef)
+                if _is_numpy_namespace(xp):
+                    weights = coef.reshape((self.base_loss.n_classes, -1), order="F")
+                else:
+                    weights = xp.reshape(coef, (-1, self.base_loss.n_classes)).T
             else:
                 weights = coef
             if self.fit_intercept:
@@ -351,6 +357,11 @@ def loss_gradient(
 
         grad_pointwise /= sw_sum
 
+        # TODO: The lbfgs solver currently only works with numpy arrays and expects
+        # both the coef argument and the result to be in numpy. So, for the time
+        # being, we stick with a mix of numpy and array API. We could consider
+        # moving coef to the device of X if lbgfs also becomes compatible with the
+        # array API.
         if not self.base_loss.is_multiclass:
             grad = np.empty_like(coef, dtype=weights.dtype)
             X_grad = X.T @ grad_pointwise
@@ -431,23 +442,30 @@ def gradient(
             sample_weight=sample_weight,
             n_threads=n_threads,
         )
-        sw_sum = n_samples if sample_weight is None else np.sum(sample_weight)
+        xp, _, device = get_namespace_and_device(X, y, sample_weight)
+        sw_sum = n_samples if sample_weight is None else xp.sum(sample_weight)
         grad_pointwise /= sw_sum
 
         if not self.base_loss.is_multiclass:
-            grad = np.empty_like(coef, dtype=weights.dtype)
+            grad = xp.empty_like(coef, dtype=weights.dtype, device=device)
             grad[:n_features] = X.T @ grad_pointwise + l2_reg_strength * weights
             if self.fit_intercept:
-                grad[-1] = grad_pointwise.sum()
+                grad[-1] = xp.sum(grad_pointwise)
             return grad
         else:
-            grad = np.empty((n_classes, n_dof), dtype=weights.dtype, order="F")
+            if _is_numpy_namespace(xp):
+                grad = np.empty((n_classes, n_dof), dtype=weights.dtype, order="F")
+            else:
+                grad = xp.empty((n_classes, n_dof), dtype=weights.dtype, device=device)
             # gradient.shape = (n_samples, n_classes)
             grad[:, :n_features] = grad_pointwise.T @ X + l2_reg_strength * weights
             if self.fit_intercept:
-                grad[:, -1] = grad_pointwise.sum(axis=0)
+                grad[:, -1] = xp.sum(grad_pointwise, axis=0)
             if coef.ndim == 1:
-                return grad.ravel(order="F")
+                if _is_numpy_namespace(xp):
+                    return grad.ravel(order="F")
+                else:
+                    return _ravel(grad.T, xp=xp)
             else:
                 return grad
 
@@ -756,7 +774,9 @@ def gradient_hessian_product(
         (n_samples, n_features), n_classes = X.shape, self.base_loss.n_classes
         n_dof = n_features + int(self.fit_intercept)
         weights, intercept, raw_prediction = self.weight_intercept_raw(coef, X)
-        sw_sum = n_samples if sample_weight is None else np.sum(sample_weight)
+        xp, _, device = get_namespace_and_device(X, y, sample_weight)
+        is_numpy_ns = _is_numpy_namespace(xp)
+        sw_sum = n_samples if sample_weight is None else xp.sum(sample_weight)
 
         if not self.base_loss.is_multiclass:
             grad_pointwise, hess_pointwise = self.base_loss.gradient_hessian(
@@ -767,39 +787,45 @@ def gradient_hessian_product(
             )
             grad_pointwise /= sw_sum
             hess_pointwise /= sw_sum
-            grad = np.empty_like(coef, dtype=weights.dtype)
+            grad = xp.empty_like(coef, dtype=weights.dtype)
             grad[:n_features] = X.T @ grad_pointwise + l2_reg_strength * weights
             if self.fit_intercept:
-                grad[-1] = grad_pointwise.sum()
+                grad[-1] = xp.sum(grad_pointwise)
 
             # Precompute as much as possible: hX, hX_sum and hessian_sum
-            hessian_sum = hess_pointwise.sum()
+            hessian_sum = xp.sum(hess_pointwise)
             if sparse.issparse(X):
                 hX = (
                     sparse.dia_array((hess_pointwise, 0), shape=(n_samples, n_samples))
                     @ X
                 )
             else:
-                hX = hess_pointwise[:, np.newaxis] * X
+                hX = hess_pointwise[:, None] * X
 
             if self.fit_intercept:
-                # Calculate the double derivative with respect to intercept.
-                # Note: In case hX is sparse, hX.sum is a matrix object.
-                hX_sum = np.squeeze(np.asarray(hX.sum(axis=0)))
-                # prevent squeezing to zero-dim array if n_features == 1
-                hX_sum = np.atleast_1d(hX_sum)
+                # Calculate the double derivative with respect to the intercept.
+                if sparse.issparse(X):
+                    # Note: In case hX is sparse, hX.sum is a matrix object.
+                    hX_sum = xp.asarray(hX.sum(axis=0))
+                else:
+                    hX_sum = xp.sum(hX, axis=0)
 
             # With intercept included and l2_reg_strength = 0, hessp returns
             # res = (X, 1)' @ diag(h) @ (X, 1) @ s
             #     = (X, 1)' @ (hX @ s[:n_features], sum(h) * s[-1])
             # res[:n_features] = X' @ hX @ s[:n_features] + sum(h) * s[-1]
             # res[-1] = 1' @ hX @ s[:n_features] + sum(h) * s[-1]
             def hessp(s):
-                ret = np.empty_like(s)
+                ret = xp.empty_like(s)
                 if sparse.issparse(X):
                     ret[:n_features] = X.T @ (hX @ s[:n_features])
                 else:
-                    ret[:n_features] = np.linalg.multi_dot([X.T, hX, s[:n_features]])
+                    if is_numpy_ns:
+                        ret[:n_features] = np.linalg.multi_dot(
+                            [X.T, hX, s[:n_features]]
+                        )
+                    else:
+                        ret[:n_features] = X.T @ (hX @ s[:n_features])
                 ret[:n_features] += l2_reg_strength * s[:n_features]
 
                 if self.fit_intercept:
@@ -820,10 +846,13 @@ def hessp(s):
                 n_threads=n_threads,
             )
             grad_pointwise /= sw_sum
-            grad = np.empty((n_classes, n_dof), dtype=weights.dtype, order="F")
+            if is_numpy_ns:
+                grad = np.empty((n_classes, n_dof), dtype=weights.dtype, order="F")
+            else:
+                grad = xp.empty((n_classes, n_dof), dtype=weights.dtype, device=device)
             grad[:, :n_features] = grad_pointwise.T @ X + l2_reg_strength * weights
             if self.fit_intercept:
-                grad[:, -1] = grad_pointwise.sum(axis=0)
+                grad[:, -1] = xp.sum(grad_pointwise, axis=0)
 
             # Full hessian-vector product, i.e. not only the diagonal part of the
             # hessian. Derivation with some index battle for input vector s:
@@ -847,29 +876,47 @@ def hessp(s):
             #
             # See also https://github.com/scikit-learn/scikit-learn/pull/3646#discussion_r17461411
             def hessp(s):
-                s = s.reshape((n_classes, -1), order="F")  # shape = (n_classes, n_dof)
+                if is_numpy_ns:
+                    s = s.reshape(
+                        (n_classes, -1), order="F"
+                    )  # shape = (n_classes, n_dof)
+                else:
+                    s = xp.reshape(s, (-1, n_classes)).T
                 if self.fit_intercept:
                     s_intercept = s[:, -1]
                     s = s[:, :-1]  # shape = (n_classes, n_features)
                 else:
                     s_intercept = 0
                 tmp = X @ s.T + s_intercept  # X_{im} * s_k_m
-                tmp -= (proba * tmp).sum(axis=1)[:, np.newaxis]  # - sum_l ..
+                tmp -= xp.sum(proba * tmp, axis=1)[:, None]  # - sum_l ..
                 tmp *= proba  # * p_i_k
                 if sample_weight is not None:
-                    tmp *= sample_weight[:, np.newaxis]
+                    tmp *= sample_weight[:, None]
                 # hess_prod = empty_like(grad), but we ravel grad below and this
                 # function is run after that.
-                hess_prod = np.empty((n_classes, n_dof), dtype=weights.dtype, order="F")
+                if is_numpy_ns:
+                    hess_prod = np.empty(
+                        (n_classes, n_dof), dtype=weights.dtype, order="F"
+                    )
+                else:
+                    hess_prod = xp.empty(
+                        (n_classes, n_dof), dtype=weights.dtype, device=device
+                    )
                 hess_prod[:, :n_features] = (tmp.T @ X) / sw_sum + l2_reg_strength * s
                 if self.fit_intercept:
-                    hess_prod[:, -1] = tmp.sum(axis=0) / sw_sum
+                    hess_prod[:, -1] = xp.sum(tmp, axis=0) / sw_sum
                 if coef.ndim == 1:
-                    return hess_prod.ravel(order="F")
+                    if is_numpy_ns:
+                        return hess_prod.ravel(order="F")
+                    else:
+                        return _ravel(hess_prod.T, xp=xp)
                 else:
                     return hess_prod
 
             if coef.ndim == 1:
-                return grad.ravel(order="F"), hessp
+                if is_numpy_ns:
+                    return grad.ravel(order="F"), hessp
+                else:
+                    return _ravel(grad.T, xp=xp), hessp
 
         return grad, hessp