scikit-learn · jjerphan · Feb 23, 2022 · Feb 23, 2022 · Feb 24, 2022 · Feb 25, 2022
diff --git a/.gitignore b/.gitignore
@@ -87,3 +87,4 @@ sklearn/utils/_weight_vector.pxd
 sklearn/linear_model/_sag_fast.pyx
 sklearn/metrics/_dist_metrics.pyx
 sklearn/metrics/_dist_metrics.pxd
+sklearn/metrics/_pairwise_distances_reduction.pyx
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
@@ -186,6 +186,9 @@ Changelog
 
 
 - |Efficiency| Low-level routines for reductions on pairwise distances
+  for dense float32 and float64 datasets have been refactored.
+  The following functions and estimators now benefit from improved performances,
+  in particular on multi-cores machines:
   for dense float64 datasets have been refactored. The following functions
   and estimators now benefit from improved performances in terms of hardware
   scalability and speed-ups:

diff --git a/sklearn/metrics/_dist_metrics.pxd.tp b/sklearn/metrics/_dist_metrics.pxd.tp
@@ -100,12 +100,11 @@ cdef class DistanceMetric{{name_suffix}}:
 
     cdef DTYPE_t _dist_to_rdist(self, {{DTYPE_t}} dist) nogil except -1
 
-{{endfor}}
 
 ######################################################################
 # DatasetsPair base class
-cdef class DatasetsPair:
-    cdef DistanceMetric distance_metric
+cdef class DatasetsPair{{name_suffix}}:
+    cdef DistanceMetric{{name_suffix}} distance_metric
 
     cdef ITYPE_t n_samples_X(self) nogil
 
@@ -116,8 +115,10 @@ cdef class DatasetsPair:
     cdef DTYPE_t surrogate_dist(self, ITYPE_t i, ITYPE_t j) nogil
 
 
-cdef class DenseDenseDatasetsPair(DatasetsPair):
+cdef class DenseDenseDatasetsPair{{name_suffix}}(DatasetsPair{{name_suffix}}):
     cdef:
-        const DTYPE_t[:, ::1] X
-        const DTYPE_t[:, ::1] Y
+        const {{DTYPE_t}}[:, ::1] X
+        const {{DTYPE_t}}[:, ::1] Y
         ITYPE_t d
+
+{{endfor}}
diff --git a/sklearn/metrics/_dist_metrics.pyx.tp b/sklearn/metrics/_dist_metrics.pyx.tp
@@ -1170,11 +1170,10 @@ cdef class PyFuncDistance{{name_suffix}}(DistanceMetric{{name_suffix}}):
             raise TypeError("Custom distance function must accept two "
                             "vectors and return a float.")
 
-{{endfor}}
 
 ######################################################################
 # Datasets Pair Classes
-cdef class DatasetsPair:
+cdef class DatasetsPair{{name_suffix}}:
     """Abstract class which wraps a pair of datasets (X, Y).
 
     This class allows computing distances between a single pair of rows of
@@ -1211,7 +1210,7 @@ cdef class DatasetsPair:
         Y,
         str metric="euclidean",
         dict metric_kwargs=None,
-    ) -> DatasetsPair:
+    ) -> DatasetsPair{{name_suffix}}:
         """Return the DatasetsPair implementation for the given arguments.
 
         Parameters
@@ -1241,14 +1240,14 @@ cdef class DatasetsPair:
             The suited DatasetsPair implementation.
         """
         cdef:
-            DistanceMetric distance_metric = DistanceMetric.get_metric(
+            DistanceMetric{{name_suffix}} distance_metric = DistanceMetric{{name_suffix}}.get_metric(
                 metric,
                 **(metric_kwargs or {})
             )
 
-        if not(X.dtype == Y.dtype == np.float64):
+        if not(X.dtype == Y.dtype and X.dtype in DatasetsPair{{name_suffix}}.valid_dtypes()):
             raise ValueError(
-                f"Only 64bit float datasets are supported at this time, "
+                f"Only np.float64 and np.float32 datasets are supported at this time, "
                 f"got: X.dtype={X.dtype} and Y.dtype={Y.dtype}."
             )
 
@@ -1260,11 +1259,15 @@ cdef class DatasetsPair:
         if issparse(X) or issparse(Y):
             raise ValueError("Only dense datasets are supported for X and Y.")
 
-        return DenseDenseDatasetsPair(X, Y, distance_metric)
+        return DenseDenseDatasetsPair{{name_suffix}}(X, Y, distance_metric)
 
-    def __init__(self, DistanceMetric distance_metric):
+    def __init__(self, DistanceMetric{{name_suffix}} distance_metric):
         self.distance_metric = distance_metric
 
+    @classmethod
+    def valid_dtypes(cls):
+        return (np.float64, np.float32)
+
     cdef ITYPE_t n_samples_X(self) nogil:
         """Number of samples in X."""
         # This is a abstract method.
@@ -1289,7 +1292,7 @@ cdef class DatasetsPair:
         return -1
 
 @final
-cdef class DenseDenseDatasetsPair(DatasetsPair):
+cdef class DenseDenseDatasetsPair{{name_suffix}}(DatasetsPair{{name_suffix}}):
     """Compute distances between row vectors of two arrays.
 
     Parameters
@@ -1305,7 +1308,7 @@ cdef class DenseDenseDatasetsPair(DatasetsPair):
         between two row vectors of (X, Y).
        """
 
-    def __init__(self, X, Y, DistanceMetric distance_metric):
+    def __init__(self, X, Y, DistanceMetric{{name_suffix}} distance_metric):
         super().__init__(distance_metric)
         # Arrays have already been checked
         self.X = X
@@ -1331,3 +1334,5 @@ cdef class DenseDenseDatasetsPair(DatasetsPair):
         return self.distance_metric.dist(&self.X[i, 0],
                                          &self.Y[j, 0],
                                          self.d)
+
+{{endfor}}