Fix sample_weight handling in KernelRidge

python/cuml/cuml/kernel_ridge/kernel_ridge.py

@@ -11,7 +11,7 @@
 
 from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.common.doc_utils import generate_docstring
-from cuml.internals import reflect
+from cuml.internals import reflect, run_in_internal_context
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base
 from cuml.internals.interop import (
@@ -21,7 +21,12 @@
     to_gpu,
 )
 from cuml.internals.mixins import RegressorMixin
-from cuml.internals.validation import check_inputs, check_is_fitted
+from cuml.internals.validation import (
+    check_consistent_length,
+    check_inputs,
+    check_is_fitted,
+    check_sample_weight,
+)
 from cuml.metrics import pairwise_kernels
 
 
@@ -59,8 +64,6 @@ def _solve_cholesky_kernel(K, y, alpha, sample_weight=None):
     has_sw = sample_weight is not None
 
     if has_sw:
-        # Unlike other solvers, we need to support sample_weight directly
-        # because K might be a pre-computed kernel.
         sw = cp.sqrt(cp.atleast_1d(sample_weight))
         y = y * sw[:, cp.newaxis]
         K *= cp.outer(sw, sw)
@@ -269,6 +272,7 @@ def __init__(
     def _more_static_tags():
         return {"multioutput": True}
 
+    @run_in_internal_context
     def _get_kernel(self, X, Y=None):
         if isinstance(self.kernel, str):
             params = {
@@ -287,11 +291,10 @@ def _get_kernel(self, X, Y=None):
     def fit(
         self, X, y, sample_weight=None, *, convert_dtype=True
     ) -> "KernelRidge":
-        X, y, sample_weight, index = check_inputs(
+        X, y, index = check_inputs(
             self,
             X,
             y,
-            sample_weight,
             dtype=("float32", "float64"),
             convert_dtype=convert_dtype,
             accept_multi_output=True,
@@ -301,6 +304,14 @@ def fit(
         if ravel := (y.ndim == 1):
             y = y.reshape(-1, 1)
 
+        # Unlike other solvers, we need to special-case scalar sample weights,
+        # because K might be a pre-computed kernel.
+        if not (np.isscalar(sample_weight) and np.isfinite(sample_weight)):
+            sample_weight = check_sample_weight(
+                sample_weight, dtype=X.dtype, convert_dtype=convert_dtype
+            )
+            check_consistent_length(X, y, sample_weight)
+
         K = self._get_kernel(X)
         dual_coef = _solve_cholesky_kernel(
             K, y, cp.asarray(self.alpha), sample_weight

python/cuml/tests/test_kernel_ridge.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 import inspect
@@ -27,6 +27,9 @@ def gradient_norm(model, X, y, K, sw=None):
     else:
         sw = cp.atleast_1d(cp.array(sw, dtype=np.float64))
 
+    if y.ndim == 1:
+        y = y[:, None]
+
     X = cp.array(X, dtype=np.float64)
     y = cp.array(y, dtype=np.float64)
     K = cp.array(K, dtype=np.float64)
@@ -252,7 +255,21 @@ def estimator_array_strategy(draw):
         np.array([[2.0, 3.0], [4.0, 5.0]]),  # X_test
         np.array([0.1]),  # alpha
         None,  # sample_weight
-        np.float32,  # dtype
+        np.float64,  # dtype
+    ),
+    gamma=1.0,
+    degree=1,
+    coef0=0.0,
+)
+@example(
+    kernel_arg=("linear", {}),
+    arrays=(
+        np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]),  # X
+        np.array([[1.0, 2.0, 3.0]]).T,  # y
+        np.array([[2.0, 3.0], [4.0, 5.0]]),  # X_test
+        np.array([0.1]),  # alpha
+        0.5,  # sample_weight
+        np.float64,  # dtype
     ),
     gamma=1.0,
     degree=1,