MAINT Correctly resize buffers for upcasting

jjerphan · jjerphan · commit f0fc839adcc1 · 2022-05-29T10:24:56.000+02:00
diff --git a/sklearn/metrics/_pairwise_distances_reduction.pyx.tp b/sklearn/metrics/_pairwise_distances_reduction.pyx.tp
@@ -622,11 +622,13 @@ cdef class GEMMTermComputer{{bitness}}:
         ITYPE_t chunks_n_threads
         ITYPE_t dist_middle_terms_chunks_size
         ITYPE_t n_features
+        ITYPE_t chunk_size
 
         # Buffers for the `-2 * X_c @ Y_c.T` term computed via GEMM
         vector[vector[DTYPE_t]] dist_middle_terms_chunks
 
 {{if need_upcast}}
+        # Buffers for upcasting chunks of X and Y from 32bit to 64bit
         vector[vector[DTYPE_t]] X_c_upcast
         vector[vector[DTYPE_t]] Y_c_upcast
 {{endif}}
@@ -638,24 +640,28 @@ cdef class GEMMTermComputer{{bitness}}:
         ITYPE_t chunks_n_threads,
         ITYPE_t dist_middle_terms_chunks_size,
         ITYPE_t n_features,
+        ITYPE_t chunk_size,
     ):
         self.X = X
         self.Y = Y
         self.effective_n_threads = effective_n_threads
         self.chunks_n_threads = chunks_n_threads
         self.dist_middle_terms_chunks_size = dist_middle_terms_chunks_size
         self.n_features = n_features
+        self.chunk_size = chunk_size
 
         self.dist_middle_terms_chunks = vector[vector[DTYPE_t]](self.effective_n_threads)
 
 {{if need_upcast}}
+        # We populate the buffer for upcasting chunks of X and Y from 32bit to 64bit.
         self.X_c_upcast = vector[vector[DTYPE_t]](self.effective_n_threads)
         self.Y_c_upcast = vector[vector[DTYPE_t]](self.effective_n_threads)
 
-        # Buffers for upcasting chunks of X and Y from 32bit to 64bit.
+        upcast_buffer_n_elements = self.chunk_size * n_features
+
         for thread_num in range(self.effective_n_threads):
-            self.X_c_upcast[thread_num].resize(self.dist_middle_terms_chunks_size)
-            self.Y_c_upcast[thread_num].resize(self.dist_middle_terms_chunks_size)
+            self.X_c_upcast[thread_num].resize(upcast_buffer_n_elements)
+            self.Y_c_upcast[thread_num].resize(upcast_buffer_n_elements)
 {{endif}}
 
 
@@ -1556,7 +1562,8 @@ cdef class FastEuclideanPairwiseDistancesArgKmin{{bitness}}(PairwiseDistancesArg
             self.effective_n_threads,
             self.chunks_n_threads,
             dist_middle_terms_chunks_size,
-            n_features=datasets_pair.X.shape[1]
+            n_features=datasets_pair.X.shape[1],
+            chunk_size=self.chunk_size,
         )
 
         if metric_kwargs is not None and "Y_norm_squared" in metric_kwargs:
@@ -2171,7 +2178,8 @@ cdef class FastEuclideanPairwiseDistancesRadiusNeighborhood{{bitness}}(PairwiseD
             self.effective_n_threads,
             self.chunks_n_threads,
             dist_middle_terms_chunks_size,
-            n_features=datasets_pair.X.shape[1]
+            n_features=datasets_pair.X.shape[1],
+            chunk_size=self.chunk_size,
         )
 
         if metric_kwargs is not None and "Y_norm_squared" in metric_kwargs:
