ENH perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict - #34277
Merged
lorentzenchr merged 3 commits intoJun 29, 2026
Conversation
Three independent speedups found in a perf review: 1. nan_euclidean_distances: compute the present-feature-count matrix with a float matmul instead of an integer one. NumPy has no BLAS path for integer matmul, so the boolean masks are cast to the working dtype first. Counts are small integers, exactly representable, so the result is unchanged. Speeds up KNNImputer and pairwise_distances/NearestNeighbors with metric="nan_euclidean". 2. GaussianMixture (tied covariance) _estimate_log_gaussian_prob: replace the per-component loop with a single BLAS expansion of the squared distance ||Xp - mu_proj||**2, the same form the diag and spherical branches already use. Not bit-identical to the loop, but a full tied fit is unchanged in practice (identical means_, lower_bound_ within ~1e-15). 3. KNeighborsRegressor.predict: fold the per-output distance-weighting loop into a single einsum that gathers the neighbor targets once. Output identical; speeds up many-output multi-target regression, neutral for single-output. Benchmarks (median, baseline vs branch, real APIs): nan_euclidean_distances (n=10000, ny=2000): d=50 3.4x, d=200 7.8x, d=500 10.7x KNNImputer.fit_transform (15% missing): 1.20-1.25x GaussianMixture(tied).fit: 1.9x (c=5) to 5.5x (c=40 / d=100) KNeighborsRegressor.predict: 1.0x (1 output) to 2.1x (1000 outputs)
rth
force-pushed
the
perf/vectorize-numeric-hotpaths
branch
from
June 12, 2026 17:34
2f98617 to
c520ecc
Compare
lorentzenchr
approved these changes
Jun 22, 2026
lorentzenchr
left a comment
Member
There was a problem hiding this comment.
@rth Thanks yor thos nice little PR. Could you add a whatsnew entry?
PS: a pitty we missed each other in Paris.
Member
|
@rth : this is a really nice PR. I'd love to merge it. Could you add a changelog entry, that way I can merge. Thanks! |
Member
Author
|
Thanks for the review @lorentzenchr and @GaelVaroquaux ! Added a changelog. So it's a changelog entry per module if I understand correctly how it works? @lorentzenchr yes, a shame we didn't get to meet. I was there only the last day. Next time) |
prady0t
pushed a commit
to prady0t/scikit-learn
that referenced
this pull request
Sep 2, 2026
…xture and KNeighborsRegressor.predict (scikit-learn#34277)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three independent, small numerical optimizations:
nan_euclidean_distances: use a float matmul instead of an integer one (so it can use BLAS)Some benchmarks for
nan_euclidean_distances(X, Y),Xisn x d,Yfixed at 2000 rows:bench_nan_euclidean_distances.py
This is in particular used inside
KNNImputer.fit_transformwhich shows more modest gains. Benchmarks are with 15% of missing valuesbench_knn_imputer.py
GaussianMixture(covariance_type="tied").fitreplace a for loop with a vectorized expression.bench_gaussian_mixture_tied.py
KNeighborsRegressor(n_neighbors, weights="distance")bench_kneighbors_regressor.py
where n_output is the number of output in multi-output regression (so this only impacts where multiple variables are regressed at once)
I confirm that are sufficient test coverage for each changed branch. If we voluntarily introduce small calculation mistakes in the changed lines, tests are failing.