Add array API tests for pairwise_distances (scikit-learn#31658)

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>
Co-authored-by: Tim Head <betatim@gmail.com>

Author: 3 people

doc/modules/array_api.rst

@@ -156,11 +156,12 @@ Metrics
 - :func:`sklearn.metrics.pairwise.chi2_kernel`
 - :func:`sklearn.metrics.pairwise.cosine_similarity`
 - :func:`sklearn.metrics.pairwise.cosine_distances`
+- :func:`sklearn.metrics.pairwise.pairwise_distances` (only supports "cosine", "euclidean" and "l2" metrics)
 - :func:`sklearn.metrics.pairwise.euclidean_distances` (see :ref:`device_support_for_float64`)
 - :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.pairwise_kernels` (supports all `sklearn.pairwise.PAIRWISE_KERNEL_FUNCTIONS` 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`

doc/whats_new/upcoming_changes/sklearn.metrics/29822.enhancement.rst renamed to doc/whats_new/upcoming_changes/array-api/29822.enhancement.rst

@@ -2,3 +2,8 @@
   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>`.
+
+- :func:`metrics.pairwise.pairwise_distances` now supports Array API
+  compatible inputs, when the underlying `metric` does (currently
+  "cosine", "euclidean" and "l2").
+  By :user:`Emily Chen <EmilyXinyi>` and :user:`Lucy Liu <lucyleeow>`.

sklearn/metrics/tests/test_common.py

@@ -65,6 +65,7 @@
     linear_kernel,
     paired_cosine_distances,
     paired_euclidean_distances,
+    pairwise_distances,
     pairwise_kernels,
     polynomial_kernel,
     rbf_kernel,
@@ -2282,6 +2283,7 @@ def check_array_api_metric_pairwise(metric, array_namespace, device, dtype_name)
     roc_curve: [
         check_array_api_binary_classification_metric,
     ],
+    pairwise_distances: [check_array_api_metric_pairwise],
 }
 
 

sklearn/metrics/tests/test_pairwise.py

@@ -151,6 +151,45 @@ def test_pairwise_distances_for_dense_data(global_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("metric", ["cosine", "euclidean"])
+def test_pairwise_distances_array_api(array_namespace, device, dtype_name, metric):
+    # Test array API support in pairwise_distances.
+    xp = _array_api_for_tests(array_namespace, device)
+
+    rng = np.random.RandomState(0)
+    # Euclidean distance should be equivalent to calling the function.
+    X_np = rng.random_sample((5, 4)).astype(dtype_name, copy=False)
+    Y_np = rng.random_sample((5, 4)).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
+        D_xp = pairwise_distances(X_xp, metric=metric)
+        D_xp_np = _convert_to_numpy(D_xp, xp=xp)
+        assert get_namespace(D_xp)[0].__name__ == xp.__name__
+        assert D_xp.device == X_xp.device
+        assert D_xp.dtype == X_xp.dtype
+
+        D_np = pairwise_distances(X_np, metric=metric)
+        assert_allclose(D_xp_np, D_np)
+
+        # Test with Y=Y_np/Y_xp
+        D_xp = pairwise_distances(X_xp, Y=Y_xp, metric=metric)
+        D_xp_np = _convert_to_numpy(D_xp, xp=xp)
+        assert get_namespace(D_xp)[0].__name__ == xp.__name__
+        assert D_xp.device == X_xp.device
+        assert D_xp.dtype == X_xp.dtype
+
+        D_np = pairwise_distances(X_np, Y=Y_np, metric=metric)
+        assert_allclose(D_xp_np, D_np)
+
+
 @pytest.mark.parametrize("coo_container", COO_CONTAINERS)
 @pytest.mark.parametrize("csc_container", CSC_CONTAINERS)
 @pytest.mark.parametrize("bsr_container", BSR_CONTAINERS)