-
Notifications
You must be signed in to change notification settings - Fork 662
[BUG] KMeans.transform returns squared distances instead of Euclidean distances #8536
Description
Describe the bug
cuml.cluster.KMeans.transform() returns squared Euclidean distances to cluster centers, while sklearn.cluster.KMeans.transform() returns Euclidean distances.
For the first transformed sample, the distances to the two centers should be sqrt(2) and sqrt(242). cuML instead returns 2 and 242, respectively. Every cuML output value is the square of the corresponding scikit-learn distance.
Steps/Code to reproduce bug
cuML reproducer:
import numpy as np from cuml.cluster import KMeans model = KMeans(n_clusters=2,init=np.array([[1., 1.],[11., 11.],]),n_init=1,).fit(np.array([[0., 0.],[1., 1.],[2., 2.],[10., 10.],[11., 11.],[12., 12.],])) print(model.transform(np.array([[0., 0.],[10., 10.],])))
Output:
[[ 2. 242.]
[162. 2.]]
For comparison, the equivalent scikit-learn code:
import numpy as np from sklearn.cluster import KMeans model = KMeans(n_clusters=2,init=np.array([[1., 1.],[11., 11.],]),n_init=1,).fit(np.array([[0., 0.],[1., 1.],[2., 2.],[10., 10.],[11., 11.],[12., 12.],])) print(model.transform(np.array([[0., 0.],[10., 10.],])))
Output:
[[ 1.41421356 15.55634919]
[12.72792206 1.41421356]]
Expected behavior
KMeans.transform() should return the Euclidean distance from each sample to each cluster center, matching the established scikit-learn cluster-distance-space API.
For the fitted centers [1, 1] and [11, 11], the expected transformed matrix is:
[[sqrt(2), sqrt(242)],
[sqrt(162), sqrt(2) ]]
which evaluates to:
[[ 1.41421356 15.55634919]
[12.72792206 1.41421356]]
Environment details (please complete the following information):
- Environment location: Docker
- Linux Distro/Architecture: Ubuntu 24.04 / x86_64
- GPU Model/Driver: NVIDIA GeForce RTX 4090 / 595.71.05
- CUDA: 13.2
- Method of cuDF & cuML install: conda
conda list:
conda list
# packages in environment at /opt/conda/envs/rapids-26.08:
#
# Name Version Build Channel
# Name Version Build Channel
python 3.14.6 h242f9ac_102_cp314 conda-forge
numpy 2.4.6 py314h2b28147_0 conda-forge
scipy 1.16.3 py314hf07bd8e_2 conda-forge
scikit-learn 1.9.0 np2py314hf09ca88_0 conda-forge
rapids 26.08.00 cuda13_260806_c2656556 rapidsai
cuml 26.08.00 cuda13_cp311_abi3_260805_265b9da6 rapidsai
libcuml 26.08.00 cuda13_260805_265b9da6 rapidsai
cudf 26.08.00 cuda13_cp311_abi3_260805_ff5b362d rapidsai
libraft 26.08.00 cuda13_260805_ebf92684 rapidsai
libraft-headers 26.08.00 cuda13_260805_ebf92684 rapidsai
pylibraft 26.08.00 cuda13_cp311_abi3_260805_ebf92684 rapidsai
cuvs 26.08.01 cuda13_cp311_abi3_260806_25b1be43 rapidsai
libcuvs 26.08.01 cuda13_260806_25b1be43 rapidsai
cupy 14.1.1 py314hdea9c46_0 conda-forge
cupy-core 14.1.1 py314hcd3b49b_0 conda-forge
numba 0.64.0 py314h8169c2f_0 conda-forge
numba-cuda 0.30.4 py314h42812f9_0 conda-forge
rmm 26.08.00 cuda13_cp311_abi3_260805_42d059f1 rapidsai
librmm 26.08.00 cuda13_260805_42d059f1 rapidsai
cuda-version 13.3 hcbadf70_3 conda-forge
cuda-bindings 13.3.1 py314h42812f9_1 conda-forge
cuda-cudart 13.3.29 hecca717_0 conda-forge
cuda-nvrtc 13.3.33 hecca717_0 conda-forge
libcublas 13.6.0.2 h676940d_0 conda-forge
libcusolver 12.2.6.9 h676940d_0 conda-forge
libcusparse 12.8.2.51 hecca717_0 conda-forge
libcurand 10.4.3.29 h676940d_0 conda-forge
Additional context
The model is deterministic because the initial centers are provided explicitly and n_init=1. Fitting leaves the two centers at [1, 1] and [11, 11].
The relationship between the outputs is exact:
1.41421356 ** 2 = 2
15.55634919 ** 2 = 242
12.72792206 ** 2 = 162
This suggests that cuML exposes the internal squared-distance matrix without applying the final square root required by the scikit-learn-compatible transform() contract.
The cuML KMeans documentation describes this method as transforming input into cluster-distance space and refers users to scikit-learn's KMeans API. Returning squared distances changes magnitudes and can silently affect downstream estimators that consume transformed features.