Validate precomputed pairwise kernel shapes

python/cuml/cuml/metrics/pairwise_kernels.py

@@ -206,12 +206,19 @@ def pairwise_kernels(
             (n_samples_X, n_features)
         Array of pairwise kernels between samples, or a feature array.
         The shape of the array should be (n_samples_X, n_samples_X) if
-        metric == "precomputed" and (n_samples_X, n_features) otherwise.
+        metric == "precomputed" and Y is None. If Y is provided with
+        metric == "precomputed", X should have shape (n_queries, n_indexed),
+        where n_indexed must equal Y.shape[0]. Otherwise, X should have shape
+        (n_samples_X, n_features).
         Acceptable formats: cuDF DataFrame, NumPy ndarray, Numba device
         ndarray, cuda array interface compliant array like CuPy.
     Y : array-like (device or host) of shape (n_samples_Y, n_features), \
         default=None
-        A second feature array only if X has shape (n_samples_X, n_features).
+        For metrics other than "precomputed", a second feature array only if X
+        has shape (n_samples_X, n_features). For metric == "precomputed", Y
+        can be any 2D array; only Y.shape[0] is used to validate the shape of X,
+        and the values and second dimension of Y do not affect the returned
+        matrix.
         Acceptable formats: cuDF DataFrame, NumPy ndarray, Numba device
         ndarray, cuda array interface compliant array like CuPy.
     metric : str or callable (numba device function), default="linear"
@@ -245,7 +252,9 @@ def pairwise_kernels(
 
     Notes
     -----
-    If metric is 'precomputed', Y is ignored and X is returned.
+    If metric is 'precomputed', Y is only used to validate the shape of X:
+    X must be square when Y is None, or X.shape[1] must equal Y.shape[0]
+    otherwise. X is returned unchanged.
 
     Examples
     --------
@@ -283,12 +292,18 @@ def pairwise_kernels(
         Y = X
     else:
         Y = check_array(Y, input_name="Y")
-    if X.shape[1] != Y.shape[1]:
-        raise ValueError("X and Y have different dimensions.")
-
     if metric == "precomputed":
+        if X.shape[1] != Y.shape[0]:
+            raise ValueError(
+                "Precomputed metric requires shape "
+                "(n_queries, n_indexed). "
+                f"Got {X.shape} for {Y.shape[0]} indexed."
+            )
         return X
 
+    if X.shape[1] != Y.shape[1]:
+        raise ValueError("X and Y have different dimensions.")
+
     if metric in PAIRWISE_KERNEL_FUNCTIONS:
         kwds = _filter_params(
             PAIRWISE_KERNEL_FUNCTIONS[metric], filter_params, **kwds

python/cuml/tests/test_kernel_ridge.py

@@ -107,6 +107,41 @@ def bad_numba_kernel2(x, y, z):
     assert np.allclose(X, pairwise_kernels(X, metric="precomputed"))
 
 
+@pytest.mark.parametrize("shape", [(5, 3), (3, 5)])
+def test_pairwise_kernels_precomputed_requires_square(shape):
+    X = np.ones(shape)
+
+    with pytest.raises(ValueError, match="Precomputed metric requires shape"):
+        pairwise_kernels(X, metric="precomputed")
+
+
+def test_pairwise_kernels_precomputed_cross_kernel():
+    X = np.arange(6).reshape(3, 2)
+    Y = np.ones((2, 10))
+
+    result = pairwise_kernels(X, Y, metric="precomputed")
+
+    cp.testing.assert_array_equal(result, cp.asarray(X))
+
+
+def test_pairwise_kernels_precomputed_wrong_indexed_count():
+    X = np.ones((3, 4))
+    Y = np.ones((2, 10))
+
+    with pytest.raises(ValueError, match="Precomputed metric requires shape"):
+        pairwise_kernels(X, Y, metric="precomputed")
+
+
+def test_pairwise_kernels_rejects_mismatched_feature_dimensions():
+    X = np.ones((3, 4))
+    Y = np.ones((2, 5))
+
+    with pytest.raises(
+        ValueError, match=r"X and Y have different dimensions\."
+    ):
+        pairwise_kernels(X, Y, metric="linear")
+
+
 @cuda.jit(device=True)
 def custom_kernel(x, y, custom_arg=5.0):
     sum = 0.0