Test and document pairwise_distances_argmin on sparse data

scikit-learn · Sep 12, 2022 · 9723115 · 9723115
1 parent 3e357a2
commit 9723115
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sklearn/metrics/pairwise.py b/sklearn/metrics/pairwise.py
@@ -721,10 +721,10 @@ def pairwise_distances_argmin(X, Y, *, axis=1, metric="euclidean", metric_kwargs
 
     Parameters
     ----------
-    X : array-like of shape (n_samples_X, n_features)
+    X : {array-like, sparse matrix} of shape (n_samples_X, n_features)
         Array containing points.
 
-    Y : array-like of shape (n_samples_Y, n_features)
+    Y : {array-like, sparse matrix} of shape (n_samples_Y, n_features)
         Arrays containing points.
 
     axis : int, default=1

diff --git a/sklearn/metrics/tests/test_pairwise.py b/sklearn/metrics/tests/test_pairwise.py
@@ -446,7 +446,9 @@ def test_pairwise_distances_argmin_min(dtype):
     assert_array_almost_equal(vals, expected_vals)
     # sparse matrix case
     idxsp, valssp = pairwise_distances_argmin_min(Xsp, Ysp, metric="euclidean")
+    idxsp2 = pairwise_distances_argmin(Xsp, Ysp, metric="euclidean")
     assert_array_almost_equal(idxsp, expected_idx)
+    assert_array_almost_equal(idxsp2, expected_idx)
     assert_array_almost_equal(valssp, expected_vals)
     # We don't want np.matrix here
     assert type(idxsp) == np.ndarray
@@ -478,7 +480,9 @@ def test_pairwise_distances_argmin_min(dtype):
     assert_array_almost_equal(vals, expected_vals)
     # sparse matrix case
     idxsp, valssp = pairwise_distances_argmin_min(Xsp, Ysp, metric="manhattan")
+    idxsp2 = pairwise_distances_argmin(Xsp, Ysp, metric="manhattan")
     assert_array_almost_equal(idxsp, expected_idx)
+    assert_array_almost_equal(idxsp2, expected_idx)
     assert_array_almost_equal(valssp, expected_vals)
 
     # Non-euclidean Scipy distance (callable)