Return Euclidean distances from KMeans transform - #8539

Open

Hashim1999164 wants to merge 2 commits into

NVIDIA:main from

Hashim1999164:fix/kmeans-transform-euclidean

Open

Return Euclidean distances from KMeans transform #8539
Hashim1999164 wants to merge 2 commits into
NVIDIA:main from
Hashim1999164:fix/kmeans-transform-euclidean

Conversation

@Hashim1999164

@Hashim1999164 Hashim1999164 commented Sep 1, 2026 •

edited

Loading

Copy link

Copy Markdown

Summary

KMeans.transform currently returns squared distances. sklearn returns Euclidean distances.

This change takes the square root of the transform output so the Python API matches sklearn. score is unchanged.

Fixes #8536

Test plan

Added a regression test against the six point example from the issue.

@Hashim1999164


 Return Euclidean distances from KMeans transform

c1afc7b

@Hashim1999164

Hashim1999164 requested a review from a team as a code owner

September 1, 2026 16:14

@Hashim1999164

Hashim1999164 requested a review from betatim

September 1, 2026 16:14

@copy-pr-bot

copy-pr-bot Bot commented Sep 1, 2026

Copy link

Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot assigned Hashim1999164

Sep 1, 2026

@github-actions github-actions Bot added the Cython / Python label

Sep 1, 2026

@coderabbitai

coderabbitai Bot commented Sep 1, 2026 •

edited

Loading

Copy link

Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

i️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4fd9e444-7419-44d5-b83d-ce4fde2f221b

📥 Commits

Reviewing files that changed from the base of the PR and between c1afc7b and f971d8f.

📒 Files selected for processing (1)

python/cuml/cuml/cluster/kmeans.pyx

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Summary

Summary by CodeRabbit

Bug Fixes
- Corrected KMeans transformations to return Euclidean distances instead of squared-Euclidean distances.
- Improved consistency with scikit-learn distance results for transformed data.
- Ensured returned distances accurately reflect the geometric distance between each sample and its assigned cluster centers.
Tests
- Added regression coverage validating KMeans distance calculations against expected values and scikit-learn results.

Walkthrough

KMeans.transform now requests Euclidean distances directly from the backend. A regression test validates expected values and scikit-learn compatibility.

Changes

KMeans distance correction

Layer / File(s)	Summary
Backend distance selection and regression coverage `python/cuml/cuml/cluster/kmeans.pyx`, `python/cuml/tests/test_kmeans.py`	`KMeans.transform` uses `DistanceType.L2SqrtExpanded` and returns the backend output unchanged. The regression test checks fixed Euclidean distances against explicit values and scikit-learn.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to f971d

KMeans.transform now returns Euclidean distances compatible with scikit-learn, with regression coverage for expected results. No merge-blocking risk remains.

Suggested reviewers: jcrist

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (1 skipped: 1...	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the primary change: returning Euclidean distances from KMeans.transform.
Description check	✅ Passed	The description directly explains the KMeans.transform distance issue, the scikit-learn compatibility goal, and the regression test.
Linked Issues check	✅ Passed	The implementation changes KMeans.transform to request square-rooted L2 distances, and the regression test verifies the expected values against explicit calculations and scikit-learn. This satisfies i...
Out of Scope Changes check	✅ Passed	The changes are limited to the KMeans.transform implementation and a regression test for the linked distance behavior. No unrelated changes are present.

Full details: Docstring Coverage

Explanation

Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (1 skipped: 1 unsupported.)

Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

viclafargue

viclafargue requested changes

Sep 3, 2026

View reviewed changes

@viclafargue viclafargue left a comment

Copy link

Copy Markdown

Contributor

Choose a reason for hiding this comment

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

Thanks for opening a PR!

python/cuml/cuml/cluster/kmeans.pyx Outdated

Comment on lines +1167 to +1168

# C++/cuVS uses L2Expanded, which is squared Euclidean.

return cp.sqrt(out)

@viclafargue viclafargue Sep 3, 2026 •

edited

Loading

Copy link

Copy Markdown

Contributor

Choose a reason for hiding this comment

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

Could we use cuVS’s native square-root metric instead?

Set it after _kmeans_init_params() (params.metric = DistanceType.L2SqrtExpanded). cuVS then applies the square root in the distance epilogue, avoiding the extra CuPy kernel and output-sized allocation from cp.sqrt(out).

@Hashim1999164 Hashim1999164 Sep 4, 2026

Copy link

Copy Markdown

Author

Choose a reason for hiding this comment

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

Done. transform now sets params.metric to DistanceType.L2SqrtExpanded after _kmeans_init_params so cuVS applies the square root in the distance epilogue. The extra CuPy sqrt and output sized allocation are gone.

@viclafargue viclafargue added bug non-breaking labels

Sep 4, 2026

@Hashim1999164


 Use cuVS L2SqrtExpanded for KMeans transform distances.

f971d8f

Labels

bug Cython / Python non-breaking

3 participants

@Hashim1999164 @viclafargue @GPUtester

Navigation Menu

Uh oh!

Return Euclidean distances from KMeans transform - #8539

Return Euclidean distances from KMeans transform #8539
Hashim1999164 wants to merge 2 commits into
NVIDIA:main from
Hashim1999164:fix/kmeans-transform-euclidean

Conversation

@Hashim1999164 Hashim1999164 commented Sep 1, 2026 •

edited

Loading

Uh oh!

Summary

Test plan

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

coderabbitai Bot commented Sep 1, 2026 •

edited

Loading

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

@viclafargue viclafargue left a comment

Choose a reason for hiding this comment

Uh oh!

@viclafargue viclafargue Sep 3, 2026 •

edited

Loading

Uh oh!

Choose a reason for hiding this comment

Uh oh!

@Hashim1999164 Hashim1999164 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

@Hashim1999164 Hashim1999164 commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

coderabbitai Bot commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

@viclafargue viclafargue left a comment

Choose a reason for hiding this comment

Uh oh!

@viclafargue viclafargue Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

@Hashim1999164 Hashim1999164 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

@Hashim1999164 Hashim1999164 commented Sep 1, 2026 •

edited

Loading

coderabbitai Bot commented Sep 1, 2026 •

edited

Loading

@viclafargue viclafargue Sep 3, 2026 •

edited

Loading