Return Euclidean distances from KMeans transform

[Hashim1999164]

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.

[copy-pr-bot]

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ 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]

Thanks for opening a PR!

[viclafargue]

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]

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.