Array API support for pairwise kernels (scikit-learn#29822)

EmilyXinyi · lucyleeow · ogrisel · web-flow · commit cccf7b46ffdd · 2025-06-19T08:54:18.000+02:00
Co-authored-by: Lucy Liu &lt;jliu176@gmail.com&gt;
Co-authored-by: Olivier Grisel &lt;olivier.grisel@ensta.org&gt;
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -158,6 +158,7 @@ Metrics
 - :func:`sklearn.metrics.pairwise.linear_kernel`
 - :func:`sklearn.metrics.pairwise.paired_cosine_distances`
 - :func:`sklearn.metrics.pairwise.paired_euclidean_distances`
+- :func:`sklearn.metrics.pairwise.pairwise_kernels` (supports all metrics except :func:`sklearn.metrics.pairwise.laplacian_kernel`)
 - :func:`sklearn.metrics.pairwise.polynomial_kernel`
 - :func:`sklearn.metrics.pairwise.rbf_kernel` (see :ref:`device_support_for_float64`)
 - :func:`sklearn.metrics.pairwise.sigmoid_kernel`
diff --git a/doc/whats_new/upcoming_changes/sklearn.metrics/29822.enhancement.rst b/doc/whats_new/upcoming_changes/sklearn.metrics/29822.enhancement.rst
@@ -0,0 +1,4 @@
+- :func:`metrics.pairwise.pairwise_kernels` now supports Array API
+  compatible inputs, when the underling `metric` does (the only metric NOT currently
+  supported is :func:`sklearn.metrics.pairwise.laplacian_kernel`).
+  By :user:`Emily Chen <EmilyXinyi>` and :user:`Lucy Liu <lucyleeow>`.
diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py
@@ -72,6 +72,15 @@ def _return_float_dtype(X, Y):
     return X, Y, dtype
 
 
+def _find_floating_dtype_allow_sparse(X, Y, xp=None):
+    """Find matching floating type, allowing for sparse input."""
+    if any([issparse(X), issparse(Y)]) or _is_numpy_namespace(xp):
+        X, Y, dtype_float = _return_float_dtype(X, Y)
+    else:
+        dtype_float = _find_matching_floating_dtype(X, Y, xp=xp)
+    return X, Y, dtype_float
+
+
 def check_pairwise_arrays(
     X,
     Y,
@@ -177,10 +186,7 @@ def check_pairwise_arrays(
     ensure_all_finite = _deprecate_force_all_finite(force_all_finite, ensure_all_finite)
 
     xp, _ = get_namespace(X, Y)
-    if any([issparse(X), issparse(Y)]) or _is_numpy_namespace(xp):
-        X, Y, dtype_float = _return_float_dtype(X, Y)
-    else:
-        dtype_float = _find_matching_floating_dtype(X, Y, xp=xp)
+    X, Y, dtype_float = _find_floating_dtype_allow_sparse(X, Y, xp=xp)
 
     estimator = "check_pairwise_arrays"
     if dtype == "infer_float":
@@ -433,7 +439,7 @@ def _euclidean_distances(X, Y, X_norm_squared=None, Y_norm_squared=None, squared
     # Ensure that distances between vectors and themselves are set to 0.0.
     # This may not be the case due to floating point rounding errors.
     if X is Y:
-        _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False)
+        distances = _fill_or_add_to_diagonal(distances, 0, xp=xp, add_value=False)
 
     if squared:
         return distances
@@ -1171,7 +1177,7 @@ def cosine_distances(X, Y=None):
     if X is Y or Y is None:
         # Ensure that distances between vectors and themselves are set to 0.0.
         # This may not be the case due to floating point rounding errors.
-        _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False)
+        S = _fill_or_add_to_diagonal(S, 0.0, xp, add_value=False)
     return S
 
 
@@ -1943,40 +1949,48 @@ def distance_metrics():
     return PAIRWISE_DISTANCE_FUNCTIONS
 
 
-def _dist_wrapper(dist_func, dist_matrix, slice_, *args, **kwargs):
+def _transposed_dist_wrapper(dist_func, dist_matrix, slice_, *args, **kwargs):
     """Write in-place to a slice of a distance matrix."""
-    dist_matrix[:, slice_] = dist_func(*args, **kwargs)
+    dist_matrix[slice_, ...] = dist_func(*args, **kwargs).T
 
 
 def _parallel_pairwise(X, Y, func, n_jobs, **kwds):
     """Break the pairwise matrix in n_jobs even slices
     and compute them using multithreading."""
+    xp, _, device = get_namespace_and_device(X, Y)
+    X, Y, dtype_float = _find_floating_dtype_allow_sparse(X, Y, xp=xp)
 
     if Y is None:
         Y = X
-    X, Y, dtype = _return_float_dtype(X, Y)
 
     if effective_n_jobs(n_jobs) == 1:
         return func(X, Y, **kwds)
 
     # enforce a threading backend to prevent data communication overhead
-    fd = delayed(_dist_wrapper)
-    ret = np.empty((X.shape[0], Y.shape[0]), dtype=dtype, order="F")
+    fd = delayed(_transposed_dist_wrapper)
+    # Transpose `ret` such that a given thread writes its ouput to a contiguous chunk.
+    # Note `order` (i.e. F/C-contiguous) is not included in array API standard, see
+    # https://github.com/data-apis/array-api/issues/571 for details.
+    # We assume that currently (April 2025) all array API compatible namespaces
+    # allocate 2D arrays using the C-contiguity convention by default.
+    ret = xp.empty((X.shape[0], Y.shape[0]), device=device, dtype=dtype_float).T
     Parallel(backend="threading", n_jobs=n_jobs)(
-        fd(func, ret, s, X, Y[s], **kwds)
+        fd(func, ret, s, X, Y[s, ...], **kwds)
         for s in gen_even_slices(_num_samples(Y), effective_n_jobs(n_jobs))
     )
 
     if (X is Y or Y is None) and func is euclidean_distances:
         # zeroing diagonal for euclidean norm.
         # TODO: do it also for other norms.
-        np.fill_diagonal(ret, 0)
+        ret = _fill_or_add_to_diagonal(ret, 0, xp=xp, add_value=False)
 
-    return ret
+    # Transform output back
+    return ret.T
 
 
 def _pairwise_callable(X, Y, metric, ensure_all_finite=True, **kwds):
     """Handle the callable case for pairwise_{distances,kernels}."""
+    xp, _, device = get_namespace_and_device(X)
     X, Y = check_pairwise_arrays(
         X,
         Y,
@@ -1985,16 +1999,28 @@ def _pairwise_callable(X, Y, metric, ensure_all_finite=True, **kwds):
         # No input dimension checking done for custom metrics (left to user)
         ensure_2d=False,
     )
+    _, _, dtype_float = _find_floating_dtype_allow_sparse(X, Y, xp=xp)
+
+    def _get_slice(array, index):
+        # TODO: below 2 lines can be removed once min scipy >= 1.14. Support for
+        # 1D shapes in scipy sparse arrays (COO, DOK and CSR formats) only
+        # added in 1.14. We must return 2D array until min scipy 1.14.
+        if issparse(array):
+            return array[[index], :]
+        # When `metric` is a callable, 1D input arrays allowed, in which case
+        # scalar should be returned.
+        if array.ndim == 1:
+            return array[index]
+        else:
+            return array[index, ...]
 
     if X is Y:
         # Only calculate metric for upper triangle
-        out = np.zeros((X.shape[0], Y.shape[0]), dtype="float")
+        out = xp.zeros((X.shape[0], Y.shape[0]), dtype=dtype_float, device=device)
         iterator = itertools.combinations(range(X.shape[0]), 2)
         for i, j in iterator:
-            # scipy has not yet implemented 1D sparse slices; once implemented this can
-            # be removed and `arr[ind]` can be simply used.
-            x = X[[i], :] if issparse(X) else X[i]
-            y = Y[[j], :] if issparse(Y) else Y[j]
+            x = _get_slice(X, i)
+            y = _get_slice(Y, j)
             out[i, j] = metric(x, y, **kwds)
 
         # Make symmetric
@@ -2004,20 +2030,16 @@ def _pairwise_callable(X, Y, metric, ensure_all_finite=True, **kwds):
         # Calculate diagonal
         # NB: nonzero diagonals are allowed for both metrics and kernels
         for i in range(X.shape[0]):
-            # scipy has not yet implemented 1D sparse slices; once implemented this can
-            # be removed and `arr[ind]` can be simply used.
-            x = X[[i], :] if issparse(X) else X[i]
+            x = _get_slice(X, i)
             out[i, i] = metric(x, x, **kwds)
 
     else:
         # Calculate all cells
-        out = np.empty((X.shape[0], Y.shape[0]), dtype="float")
+        out = xp.empty((X.shape[0], Y.shape[0]), dtype=dtype_float)
         iterator = itertools.product(range(X.shape[0]), range(Y.shape[0]))
         for i, j in iterator:
-            # scipy has not yet implemented 1D sparse slices; once implemented this can
-            # be removed and `arr[ind]` can be simply used.
-            x = X[[i], :] if issparse(X) else X[i]
-            y = Y[[j], :] if issparse(Y) else Y[j]
+            x = _get_slice(X, i)
+            y = _get_slice(Y, j)
             out[i, j] = metric(x, y, **kwds)
 
     return out
diff --git a/sklearn/metrics/tests/test_common.py b/sklearn/metrics/tests/test_common.py
@@ -65,6 +65,7 @@
     linear_kernel,
     paired_cosine_distances,
     paired_euclidean_distances,
+    pairwise_kernels,
     polynomial_kernel,
     rbf_kernel,
     sigmoid_kernel,
@@ -2277,6 +2278,7 @@ def check_array_api_metric_pairwise(metric, array_namespace, device, dtype_name)
         check_array_api_regression_metric_multioutput,
     ],
     sigmoid_kernel: [check_array_api_metric_pairwise],
+    pairwise_kernels: [check_array_api_metric_pairwise],
     roc_curve: [
         check_array_api_binary_classification_metric,
     ],
diff --git a/sklearn/metrics/tests/test_pairwise.py b/sklearn/metrics/tests/test_pairwise.py
@@ -48,7 +48,15 @@
     sigmoid_kernel,
 )
 from sklearn.preprocessing import normalize
+from sklearn.utils._array_api import (
+    _convert_to_numpy,
+    _get_namespace_device_dtype_ids,
+    get_namespace,
+    xpx,
+    yield_namespace_device_dtype_combinations,
+)
 from sklearn.utils._testing import (
+    _array_api_for_tests,
     assert_allclose,
     assert_almost_equal,
     assert_array_equal,
@@ -295,10 +303,18 @@ def test_pairwise_precomputed_non_negative():
 
 
 def callable_rbf_kernel(x, y, **kwds):
+    xp, _ = get_namespace(x, y)
     # Callable version of pairwise.rbf_kernel.
-    K = rbf_kernel(np.atleast_2d(x), np.atleast_2d(y), **kwds)
+    K = rbf_kernel(
+        xpx.atleast_nd(x, ndim=2, xp=xp), xpx.atleast_nd(y, ndim=2, xp=xp), **kwds
+    )
     # unpack the output since this is a scalar packed in a 0-dim array
-    return K.item()
+    # Note below is array API version of numpys `item()`
+    if K.ndim > 0:
+        K_flat = xp.reshape(K, (-1,))
+        if K_flat.shape == (1,):
+            return float(K_flat[0])
+    raise ValueError("can only convert an array of size 1 to a Python scalar")
 
 
 @pytest.mark.parametrize(
@@ -334,6 +350,53 @@ def test_pairwise_parallel(func, metric, kwds, dtype):
     assert_allclose(S, S2)
 
 
+@pytest.mark.parametrize(
+    "array_namespace, device, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
+@pytest.mark.parametrize(
+    "func, metric, kwds",
+    [
+        (pairwise_distances, "euclidean", {}),
+        (pairwise_kernels, "polynomial", {"degree": 1}),
+        (pairwise_kernels, callable_rbf_kernel, {"gamma": 0.1}),
+    ],
+)
+def test_pairwise_parallel_array_api(
+    func, metric, kwds, array_namespace, device, dtype_name
+):
+    xp = _array_api_for_tests(array_namespace, device)
+    rng = np.random.RandomState(0)
+    # Why 5 and not more? this seems to still result in a lot of 0 vaules?
+    X_np = np.array(5 * rng.random_sample((5, 4)), dtype=dtype_name)
+    Y_np = np.array(5 * rng.random_sample((3, 4)), dtype=dtype_name)
+    X_xp = xp.asarray(X_np, device=device)
+    Y_xp = xp.asarray(Y_np, device=device)
+
+    with config_context(array_api_dispatch=True):
+        for y_val in (None, "not none"):
+            Y_xp = None if y_val is None else Y_xp
+            Y_np = None if y_val is None else Y_np
+
+            n_job1_xp = func(X_xp, Y_xp, metric=metric, n_jobs=1, **kwds)
+            n_job1_xp_np = _convert_to_numpy(n_job1_xp, xp=xp)
+            assert get_namespace(n_job1_xp)[0].__name__ == xp.__name__
+            assert n_job1_xp.device == X_xp.device
+            assert n_job1_xp.dtype == X_xp.dtype
+
+            n_job2_xp = func(X_xp, Y_xp, metric=metric, n_jobs=2, **kwds)
+            n_job2_xp_np = _convert_to_numpy(n_job2_xp, xp=xp)
+            assert get_namespace(n_job2_xp)[0].__name__ == xp.__name__
+            assert n_job2_xp.device == X_xp.device
+            assert n_job2_xp.dtype == X_xp.dtype
+
+            n_job2_np = func(X_np, metric=metric, n_jobs=2, **kwds)
+
+            assert_allclose(n_job1_xp_np, n_job2_xp_np)
+            assert_allclose(n_job2_xp_np, n_job2_np)
+
+
 def test_pairwise_callable_nonstrict_metric():
     # paired_distances should allow callable metric where metric(x, x) != 0
     # Knowing that the callable is a strict metric would allow the diagonal to
@@ -378,6 +441,52 @@ def test_pairwise_kernels(metric, csr_container):
     assert_allclose(K1, K2)
 
 
+@pytest.mark.parametrize(
+    "array_namespace, device, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
+@pytest.mark.parametrize(
+    "metric",
+    ["rbf", "sigmoid", "polynomial", "linear", "chi2", "additive_chi2"],
+)
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_pairwise_kernels_array_api(
+    metric, csr_container, array_namespace, device, dtype_name
+):
+    # Test array API support in pairwise_kernels.
+    xp = _array_api_for_tests(array_namespace, device)
+
+    rng = np.random.RandomState(0)
+    X_np = 10 * rng.random_sample((5, 4))
+    X_np = X_np.astype(dtype_name, copy=False)
+    Y_np = 10 * rng.random_sample((2, 4))
+    Y_np = Y_np.astype(dtype_name, copy=False)
+    X_xp = xp.asarray(X_np, device=device)
+    Y_xp = xp.asarray(Y_np, device=device)
+
+    with config_context(array_api_dispatch=True):
+        # Test with Y=None
+        K_xp = pairwise_kernels(X_xp, metric=metric)
+        K_xp_np = _convert_to_numpy(K_xp, xp=xp)
+        assert get_namespace(K_xp)[0].__name__ == xp.__name__
+        assert K_xp.device == X_xp.device
+        assert K_xp.dtype == X_xp.dtype
+
+        K_np = pairwise_kernels(X_np, metric=metric)
+        assert_allclose(K_xp_np, K_np)
+
+        # Test with Y=Y_np/Y_xp
+        K_xp = pairwise_kernels(X_xp, Y=Y_xp, metric=metric)
+        K_xp_np = _convert_to_numpy(K_xp, xp=xp)
+        assert get_namespace(K_xp)[0].__name__ == xp.__name__
+        assert K_xp.device == X_xp.device
+        assert K_xp.dtype == X_xp.dtype
+
+        K_np = pairwise_kernels(X_np, Y=Y_np, metric=metric)
+        assert_allclose(K_xp_np, K_np)
+
+
 def test_pairwise_kernels_callable():
     # Test the pairwise_kernels helper function
     # with a callable function, with given keywords.
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -557,6 +557,10 @@ def _fill_or_add_to_diagonal(array, value, xp, add_value=True, wrap=False):
         array_flat[:end:step] += value
     else:
         array_flat[:end:step] = value
+    # `array_flat` is not always a view on `array` (e.g. for certain array types that
+    # were filled via parallel processing i.e., in `_parallel_pairwise`), thus we need
+    # to return reshaped `array_flat`.
+    return xp.reshape(array_flat, array.shape)
 
 
 def _is_xp_namespace(xp, name):
