Thanks to visit codestin.com
Credit goes to github.com

Skip to content

ENH perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict - #34277

Merged
lorentzenchr merged 3 commits into
scikit-learn:mainfrom
rth:perf/vectorize-numeric-hotpaths
Jun 29, 2026
Merged

ENH perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict#34277
lorentzenchr merged 3 commits into
scikit-learn:mainfrom
rth:perf/vectorize-numeric-hotpaths

Conversation

@rth

@rth rth commented Jun 12, 2026

Copy link
Copy Markdown
Member

Three independent, small numerical optimizations:

  1. 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), X is n x d, Y fixed at 2000 rows:
bench_nan_euclidean_distances.py

n d old (s) new (s) speedup
5000 200 0.69 0.085 8.1x
10000 200 1.4 0.18 7.9x
20000 200 2.8 0.37 7.5x
10000 50 0.40 0.12 3.4x
10000 500 3.2 0.30 11x

This is in particular used inside KNNImputer.fit_transform which shows more modest gains. Benchmarks are with 15% of missing values

bench_knn_imputer.py

n d old (s) new (s) speedup
4000 100 2.64 2.11 1.25x
8000 50 5.83 4.84 1.20x
  1. GaussianMixture(covariance_type="tied").fit replace a for loop with a vectorized expression.

bench_gaussian_mixture_tied.py

n d n_components old (s) new (s) speedup
20000 50 5 0.18 0.096 1.9x
20000 50 20 0.72 0.21 3.4x
20000 50 40 1.9 0.31 6.0x
10000 50 20 0.42 0.10 4.3x
30000 50 20 1.1 0.33 3.4x
20000 30 20 0.40 0.14 2.8x
20000 100 20 1.9 0.33 5.8x
  1. KNeighborsRegressor.predict: use einsum instead of a for loop. This only applies when the estimator was initialized with KNeighborsRegressor(n_neighbors, weights="distance")

bench_kneighbors_regressor.py

n_outputs old (s) new (s) speedup
1 0.138 0.135 1.02x
10 0.150 0.145 1.03x
50 0.211 0.166 1.27x
200 0.430 0.279 1.54x
1000 1.72 0.814 2.11x

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.

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
rth force-pushed the perf/vectorize-numeric-hotpaths branch from 2f98617 to c520ecc Compare June 12, 2026 17:34

@lorentzenchr lorentzenchr left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@rth Thanks yor thos nice little PR. Could you add a whatsnew entry?

PS: a pitty we missed each other in Paris.

@GaelVaroquaux

Copy link
Copy Markdown
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!

@rth

rth commented Jun 29, 2026

Copy link
Copy Markdown
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)

@lorentzenchr lorentzenchr changed the title perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict ENH perf: numeric optimisation in nan_euclidean_distances, GaussianMixture and KNeighborsRegressor.predict Jun 29, 2026
@lorentzenchr
lorentzenchr merged commit 2ead92d into scikit-learn:main Jun 29, 2026
39 checks passed
@rth
rth deleted the perf/vectorize-numeric-hotpaths branch July 1, 2026 07:09
prady0t pushed a commit to prady0t/scikit-learn that referenced this pull request Sep 2, 2026
@jeremiedbb jeremiedbb mentioned this pull request Sep 8, 2026
14 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants