Add LSMR solver to LinearRegression

docs/source/cuml-accel/limitations.rst

@@ -246,7 +246,6 @@ LinearRegression
 ``LinearRegression`` will fall back to CPU in the following cases:
 
 - If ``positive=True``.
-- If ``X`` is sparse.
 
 Additionally, the following fitted attributes are currently not computed:
 

python/cuml/cuml/linear_model/base.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 import cupy as cp
 import cupyx.scipy.sparse
+import cupyx.scipy.sparse.linalg
+import numpy as np
 
 import cuml.internals
 from cuml.common.doc_utils import generate_docstring
@@ -192,3 +194,143 @@ def center_and_scale(
         sqrt_weight = None
 
     return X, y, X_offset, y_offset, sqrt_weight
+
+
+_ridge_transform = cp.ElementwiseKernel(
+    "T x, T s, T alpha",
+    "T out",
+    "out = s < 1e-10 ? 0 : x * s / (s * s + alpha)",
+    "_ridge_transform",
+)
+_ridge_transform_zero_alpha = cp.ElementwiseKernel(
+    "T x, T s",
+    "T out",
+    "out = s < 1e-10 ? 0 : x / s",
+    "_ridge_transform_zero_alpha",
+)
+
+
+def fit_least_squares(
+    X,
+    y,
+    sample_weight=None,
+    *,
+    fit_intercept=True,
+    alpha=0.0,
+    solver="svd",
+    tol=1e-6,
+    max_iter=None,
+    may_mutate_X=False,
+    may_mutate_y=False,
+):
+    """Fit a (possibly regularized) least-squares problem.
+
+    Parameters
+    ----------
+    X : cp.ndarray or cupyx.scipy.sparse.sp_matrix, shape (n_samples, n_features)
+        The features.
+    y : cp.ndarray array, shape (n_samples,) or (n_samples, n_targets)
+        The targets.
+    sample_weight : cp.ndarray or None
+        The sample weights.
+    fit_intercept : bool
+        Whether to fit an intercept.
+    alpha : float or cp.ndarray
+        Ridge regularization strength. Must be a non-negative float. Defaults
+        to 0 for no regularization (a LinearRegression).
+    solver : {'svd', 'lsmr'}
+        The solver to use.
+    tol : float
+        Tolerance, used by the LSMR solver.
+    max_iter : int or None
+        Maximum number of iterations, used by the LSMR solver.
+    may_mutate_X : bool
+        Whether to allow mutating X inplace to save memory when possible.
+    may_mutate_y : bool
+        Whether to allow mutating y inplace to save memory when possible.
+
+    Returns
+    -------
+    coef : cp.ndarray, shape (n_features,) or (n_targets, n_features)
+        The fitted coefficients. Returns a 1D array if y is 1D, 2D otherwise.
+    intercept : float or cp.ndarray, shape (n_targets,)
+        The intercept. A scalar if y is 1D, otherwise an array.
+    n_iter : np.ndarray or None
+        The number of solver iterations ran per-target if using the LSMR
+        solver, None otherwise.
+    """
+    y_1d = y.ndim == 1
+
+    X, y, X_offset, y_offset, sqrt_weight = center_and_scale(
+        X,
+        y,
+        sample_weight=sample_weight,
+        fit_intercept=fit_intercept,
+        may_mutate_X=may_mutate_X,
+        may_mutate_y=may_mutate_y,
+    )
+
+    # Normalize alpha to a cupy array of shape (n_targets,)
+    if cp.isscalar(alpha):
+        alpha = cp.full(y.shape[1], alpha, dtype=X.dtype)
+
+    if solver == "svd":
+        # Solve using SVD method
+        u, s, vh = cp.linalg.svd(X, full_matrices=False)
+        if (alpha == 0).all():
+            # Small optimization in the case of all-zero alpha
+            temp = _ridge_transform_zero_alpha(u.T.dot(y), s[:, None])
+        else:
+            temp = _ridge_transform(u.T.dot(y), s[:, None], alpha)
+        coef = vh.T.dot(temp).T
+        n_iter = None
+    elif solver == "lsmr":
+        if cupyx.scipy.sparse.issparse(X) and fit_intercept:
+            # To keep sparsity, sparse inputs aren't already centered when
+            # fitting an intercept. We handle removing the offset within the
+            # fit via a LinearOperator.
+            if sqrt_weight is None:
+                A = cupyx.scipy.sparse.linalg.LinearOperator(
+                    shape=X.shape,
+                    matvec=lambda w: X.dot(w) - w.dot(X_offset),
+                    rmatvec=lambda w: X.T.dot(w) - X_offset * w.sum(),
+                )
+            else:
+                A = cupyx.scipy.sparse.linalg.LinearOperator(
+                    shape=X.shape,
+                    matvec=lambda w: X.dot(w) - sqrt_weight * w.dot(X_offset),
+                    rmatvec=lambda w: X.T.dot(w)
+                    - X_offset * w.dot(sqrt_weight),
+                )
+        else:
+            A = X
+
+        coef = cp.empty((y.shape[1], X.shape[1]), dtype=X.dtype)
+        n_iter = np.empty(y.shape[1], dtype=np.int32)
+        damp = cp.sqrt(alpha)
+
+        for i in range(y.shape[1]):
+            b = y[:, i]
+            info = cupyx.scipy.sparse.linalg.lsmr(
+                A,
+                b,
+                damp=damp[i],
+                atol=tol,
+                btol=tol,
+                maxiter=max_iter,
+            )
+            coef[i] = info[0]
+            n_iter[i] = info[2]
+    else:
+        raise ValueError(f"Unsupported solver={solver!r}")
+
+    if fit_intercept:
+        intercept = y_offset - cp.dot(X_offset, coef.T)
+        if y_1d:
+            intercept = coef.dtype.type(intercept.item())
+    else:
+        intercept = 0.0
+    if y_1d:
+        coef = coef.ravel()
+
+    return coef, intercept, n_iter