From b2c8a65f4da83f603e3a4c401d11ba988adb3266 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 4 Mar 2026 15:35:01 -0800 Subject: [PATCH 1/9] add sample weights --- c/src/cluster/kmeans.cpp | 12 ++++------ cpp/include/cuvs/cluster/kmeans.hpp | 27 ++++++++++++++-------- cpp/src/cluster/detail/kmeans_balanced.cuh | 9 ++++++++ cpp/src/cluster/kmeans_cluster_cost.cu | 12 ++++++---- 4 files changed, 38 insertions(+), 22 deletions(-) diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 698f387616..622f52ff44 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -63,12 +63,10 @@ void _fit(cuvsResources_t res, RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); } else { auto kmeans_params = convert_balanced_params(params); - cuvs::cluster::kmeans::fit(*res_ptr, - kmeans_params, - cuvs::core::from_dlpack(X_tensor), - cuvs::core::from_dlpack(centroids_tensor)); - - *inertia = 0; + T inertia_temp; + auto inertia_view = raft::make_host_scalar_view(&inertia_temp); + cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, X_view, inertia_view); + *inertia = inertia_temp; *n_iter = params.hierarchical_n_iters; } } else { diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index 22c8e056ec..4ede31243b 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -1380,7 +1380,7 @@ void transform(raft::resources const& handle, raft::device_matrix_view X_new); /** - * @brief Compute cluster cost + * @brief Compute (optionally weighted) cluster cost * * @param[in] handle The raft handle * @param[in] X Training instances to cluster. The data must @@ -1390,12 +1390,16 @@ void transform(raft::resources const& handle, * row-major format. * [dim = n_clusters x n_features] * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] * */ -void cluster_cost(const raft::resources& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost); +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); /** * @brief Compute cluster cost @@ -1408,12 +1412,15 @@ void cluster_cost(const raft::resources& handle, * row-major format. * [dim = n_clusters x n_features] * @param[out] cost Resulting cluster cost - * + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] */ -void cluster_cost(const raft::resources& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost); +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); /** * @} diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 7582ec900e..1894a9a6d2 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1164,6 +1164,15 @@ void build_hierarchical(const raft::resources& handle, MathT{0.2}, mapping_op, device_memory); + + // Compute inertia if requested + if (inertia != nullptr) { + auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(cluster_centers, n_clusters, dim); + cuvs::cluster::kmeans::detail::cluster_cost( + handle, X_view, centroids_view, *inertia, std::nullopt, mapping_op); + } } } // namespace cuvs::cluster::kmeans::detail diff --git a/cpp/src/cluster/kmeans_cluster_cost.cu b/cpp/src/cluster/kmeans_cluster_cost.cu index a806cba4a4..3c877f4913 100644 --- a/cpp/src/cluster/kmeans_cluster_cost.cu +++ b/cpp/src/cluster/kmeans_cluster_cost.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,16 +11,18 @@ namespace cuvs::cluster::kmeans { void cluster_cost(const raft::resources& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::host_scalar_view cost) + raft::host_scalar_view cost, + std::optional> sample_weight) { - cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost); + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); } void cluster_cost(const raft::resources& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::host_scalar_view cost) + raft::host_scalar_view cost, + std::optional> sample_weight) { - cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost); + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); } } // namespace cuvs::cluster::kmeans From 4e8f2e47a26e75c11bf1c4f4924d495498c28959 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 4 Mar 2026 16:14:01 -0800 Subject: [PATCH 2/9] update impl --- cpp/src/cluster/detail/kmeans_balanced.cuh | 14 +++++---- cpp/src/cluster/kmeans.cuh | 35 ++++++++++++++++++---- 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 1894a9a6d2..a45493ffb9 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1165,13 +1165,15 @@ void build_hierarchical(const raft::resources& handle, mapping_op, device_memory); - // Compute inertia if requested + // Compute inertia if requested (only supported when T == MathT) if (inertia != nullptr) { - auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); - auto centroids_view = - raft::make_device_matrix_view(cluster_centers, n_clusters, dim); - cuvs::cluster::kmeans::detail::cluster_cost( - handle, X_view, centroids_view, *inertia, std::nullopt, mapping_op); + if constexpr (std::is_same_v) { + auto X_view = raft::make_device_matrix_view( + reinterpret_cast(dataset), n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(cluster_centers, n_clusters, dim); + cuvs::cluster::kmeans::detail::cluster_cost(handle, X_view, centroids_view, *inertia); + } } } diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 7d37b9cf80..633c8f346d 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -404,14 +404,27 @@ void min_cluster_distance(raft::resources const& handle, workspace); } +/** + * @brief Compute (optionally weighted) cluster cost (inertia). + * + * @tparam DataT float or double + * @tparam IndexT Index type + * + * @param[in] handle The raft handle + * @param[in] X Input data [n_samples x n_features] + * @param[in] centroids Cluster centroids [n_clusters x n_features] + * @param[out] cost Sum of squared distances to nearest centroid + * @param[in] sample_weight Optional per-sample weights [n_samples] + */ template -void cluster_cost(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost) +void cluster_cost( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt) { - auto stream = raft::resource::get_cuda_stream(handle); - + auto stream = raft::resource::get_cuda_stream(handle); auto n_clusters = centroids.extent(0); auto n_samples = X.extent(0); auto n_features = X.extent(1); @@ -440,6 +453,16 @@ void cluster_cost(raft::resources const& handle, n_clusters, workspace); + // Apply sample weights if provided + if (sample_weight.has_value()) { + raft::linalg::map( + handle, + min_dist.view(), + [] __device__(DataT d, DataT w) { return d * w; }, + raft::make_const_mdspan(min_dist.view()), + sample_weight.value()); + } + auto device_cost = raft::make_device_scalar(handle, DataT(0)); cuvs::cluster::kmeans::cluster_cost( From d6f45247ea70838fa6f779c431853a7aa51eb8d0 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 4 Mar 2026 16:20:37 -0800 Subject: [PATCH 3/9] fix min_cluster_dist --- cpp/src/cluster/kmeans.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 633c8f346d..8e26372b46 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -457,9 +457,9 @@ void cluster_cost( if (sample_weight.has_value()) { raft::linalg::map( handle, - min_dist.view(), + min_cluster_distance.view(), [] __device__(DataT d, DataT w) { return d * w; }, - raft::make_const_mdspan(min_dist.view()), + raft::make_const_mdspan(min_cluster_distance.view()), sample_weight.value()); } From 1fa9013f1534d7da131b7ae032ef55478d74d05f Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 4 Mar 2026 18:02:37 -0800 Subject: [PATCH 4/9] update instantiations --- c/src/cluster/kmeans.cpp | 4 +- cpp/include/cuvs/cluster/kmeans.hpp | 55 ++++++++++++++++++-- cpp/src/cluster/detail/kmeans_balanced.cuh | 6 ++- cpp/src/cluster/kmeans_balanced.cuh | 11 ++-- cpp/src/cluster/kmeans_balanced_fit_float.cu | 7 +-- cpp/src/cluster/kmeans_balanced_fit_half.cu | 7 +-- cpp/src/cluster/kmeans_balanced_fit_int8.cu | 7 +-- cpp/src/cluster/kmeans_balanced_fit_uint8.cu | 7 +-- cpp/src/cluster/kmeans_cluster_cost.cu | 18 +++++++ 9 files changed, 99 insertions(+), 23 deletions(-) diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 622f52ff44..57b6282c20 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -64,8 +64,8 @@ void _fit(cuvsResources_t res, } else { auto kmeans_params = convert_balanced_params(params); T inertia_temp; - auto inertia_view = raft::make_host_scalar_view(&inertia_temp); - cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, X_view, inertia_view); + auto inertia_view = raft::make_host_scalar_view(&inertia_temp); + cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, cuvs::core::from_dlpack(X_tensor), cuvs::core::from_dlpack(centroids_tensor), std::make_optional(inertia_view)); *inertia = inertia_temp; *n_iter = params.hierarchical_n_iters; } diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index 4ede31243b..ac7a9ddba7 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -11,6 +11,8 @@ #include #include +#include + namespace cuvs::cluster::kmeans { /** Base structure for parameters that are common to all k-means algorithms */ @@ -424,7 +426,8 @@ void fit(raft::resources const& handle, void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -458,7 +461,8 @@ void fit(const raft::resources& handle, void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -492,7 +496,8 @@ void fit(const raft::resources& handle, void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -526,7 +531,8 @@ void fit(const raft::resources& handle, void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Predict the closest cluster each sample in X belongs to. @@ -1422,6 +1428,47 @@ void cluster_cost( raft::host_scalar_view cost, std::optional> sample_weight = std::nullopt); +/** + * @brief Compute (optionally weighted) cluster cost + * + * @param[in] handle The raft handle + * @param[in] X Training instances to cluster. The data must + * be in row-major format. + * [dim = n_samples x n_features] + * @param[in] centroids Cluster centroids. The data must be in + * row-major format. + * [dim = n_clusters x n_features] + * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] + */ +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); + +/** + * @brief Compute (optionally weighted) cluster cost + * + * @param[in] handle The raft handle + * @param[in] X Training instances to cluster. The data must + * be in row-major format. + * [dim = n_samples x n_features] + * @param[in] centroids Cluster centroids. The data must be in + * row-major format. + * [dim = n_clusters x n_features] + * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] + */ +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); /** * @} */ diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index a45493ffb9..6fd2bfe844 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1025,7 +1025,8 @@ void build_hierarchical(const raft::resources& handle, MathT* cluster_centers, IdxT n_clusters, MappingOpT mapping_op, - const MathT* dataset_norm = nullptr) + const MathT* dataset_norm = nullptr, + MathT* inertia = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); using LabelT = uint32_t; @@ -1172,7 +1173,8 @@ void build_hierarchical(const raft::resources& handle, reinterpret_cast(dataset), n_rows, dim); auto centroids_view = raft::make_device_matrix_view(cluster_centers, n_clusters, dim); - cuvs::cluster::kmeans::detail::cluster_cost(handle, X_view, centroids_view, *inertia); + cuvs::cluster::kmeans::cluster_cost( + handle, X_view, centroids_view, raft::make_host_scalar_view(inertia)); } } } diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index d3a85b1d94..28bb9a73dd 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -67,7 +67,8 @@ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, raft::device_matrix_view centroids, - MappingOpT mapping_op = raft::identity_op()) + MappingOpT mapping_op = raft::identity_op(), + std::optional> inertia_out = std::nullopt) { RAFT_EXPECTS(X.extent(1) == centroids.extent(1), "Number of features in dataset and centroids are different"); @@ -78,6 +79,8 @@ void fit(const raft::resources& handle, "The number of centroids must be strictly positive and cannot exceed the number of " "points in the training dataset."); + MathT* inertia_ptr = inertia_out.has_value() ? inertia_out.value().data_handle() : nullptr; + cuvs::cluster::kmeans::detail::build_hierarchical(handle, params, X.extent(1), @@ -85,7 +88,9 @@ void fit(const raft::resources& handle, X.extent(0), centroids.data_handle(), centroids.extent(0), - mapping_op); + mapping_op, + static_cast(nullptr), + inertia_ptr); } /** diff --git a/cpp/src/cluster/kmeans_balanced_fit_float.cu b/cpp/src/cluster/kmeans_balanced_fit_float.cu index 37a6bb127c..f3ef94b7be 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_float.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_float.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_half.cu b/cpp/src/cluster/kmeans_balanced_fit_half.cu index e554930293..7272e6087a 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_half.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_half.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_int8.cu b/cpp/src/cluster/kmeans_balanced_fit_int8.cu index a46f7bcc94..3615c4675b 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_int8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_int8.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_uint8.cu b/cpp/src/cluster/kmeans_balanced_fit_uint8.cu index 8395dd107f..2a7211e48e 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_uint8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_uint8.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_cluster_cost.cu b/cpp/src/cluster/kmeans_cluster_cost.cu index 3c877f4913..0cdc182fb9 100644 --- a/cpp/src/cluster/kmeans_cluster_cost.cu +++ b/cpp/src/cluster/kmeans_cluster_cost.cu @@ -25,4 +25,22 @@ void cluster_cost(const raft::resources& handle, { cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); } + +void cluster_cost(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight) +{ + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); +} + +void cluster_cost(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight) +{ + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); +} } // namespace cuvs::cluster::kmeans From 4ccce8314df63db36099e7d1c2732d3f5603dc0e Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 5 Mar 2026 12:40:13 -0800 Subject: [PATCH 5/9] fix all the docs --- cpp/include/cuvs/cluster/kmeans.hpp | 8 ++++++ cpp/src/cluster/detail/kmeans_balanced.cuh | 29 +++++++++++----------- cpp/src/cluster/kmeans_balanced.cuh | 6 +++-- 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index ac7a9ddba7..a839cecf56 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -422,6 +422,8 @@ void fit(raft::resources const& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, @@ -457,6 +459,8 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, @@ -492,6 +496,8 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, @@ -527,6 +533,8 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 6fd2bfe844..cab01221e7 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -999,22 +999,23 @@ auto build_fine_clusters(const raft::resources& handle, /** * @brief Hierarchical balanced k-means * - * @tparam T element type - * @tparam MathT type of the centroids and mapped data - * @tparam IdxT index type - * @tparam LabelT label type + * @tparam T element type + * @tparam MathT type of the centroids and mapped data + * @tparam IdxT index type * @tparam MappingOpT type of the mapping operation * - * @param[in] handle The raft handle. - * @param[in] params Structure containing the hyper-parameters - * @param dim number of columns in `centers` and `dataset` - * @param[in] dataset a device pointer to the source dataset [n_rows, dim] - * @param n_rows number of rows in the input - * @param[out] cluster_centers a device pointer to the found cluster centers [n_cluster, dim] - * @param n_cluster - * @param metric the distance type - * @param mapping_op Mapping operation from T to MathT - * @param stream + * @param[in] handle The raft handle. + * @param[in] params Structure containing the hyper-parameters + * @param[in] dim Number of columns in `cluster_centers` and `dataset` + * @param[in] dataset A device pointer to the source dataset [n_rows, dim] + * @param[in] n_rows Number of rows in the input + * @param[out] cluster_centers A device pointer to the found cluster centers [n_clusters, dim] + * @param[in] n_clusters Requested number of clusters + * @param[in] mapping_op Mapping operation from T to MathT + * @param[in] dataset_norm (optional) Pre-computed L2 norms of each row in the dataset [n_rows] + * @param[out] inertia (optional) If non-null, the sum of squared distances of samples to + * their closest cluster center is written here. + * Only supported when T == MathT (float/double). */ template void build_hierarchical(const raft::resources& handle, diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index 28bb9a73dd..f1a2f39e21 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -61,6 +61,8 @@ namespace cuvs::cluster::kmeans_balanced { * @param[out] centroids The generated centroids [dim = n_clusters x n_features] * @param[in] mapping_op (optional) Functor to convert from the input datatype to the arithmetic * datatype. If DataT == MathT, this must be the identity. + * @param[out] inertia (optional) Sum of squared distances of samples to their + * closest cluster center. */ template void fit(const raft::resources& handle, @@ -68,7 +70,7 @@ void fit(const raft::resources& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, MappingOpT mapping_op = raft::identity_op(), - std::optional> inertia_out = std::nullopt) + std::optional> inertia = std::nullopt) { RAFT_EXPECTS(X.extent(1) == centroids.extent(1), "Number of features in dataset and centroids are different"); @@ -79,7 +81,7 @@ void fit(const raft::resources& handle, "The number of centroids must be strictly positive and cannot exceed the number of " "points in the training dataset."); - MathT* inertia_ptr = inertia_out.has_value() ? inertia_out.value().data_handle() : nullptr; + MathT* inertia_ptr = inertia.has_value() ? inertia.value().data_handle() : nullptr; cuvs::cluster::kmeans::detail::build_hierarchical(handle, params, From c3ca46b0d736690e5102d0f38faf083c93a46255 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 5 Mar 2026 14:17:13 -0800 Subject: [PATCH 6/9] style --- cpp/src/cluster/kmeans_balanced.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index f1a2f39e21..2b250b92cf 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -69,7 +69,7 @@ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, raft::device_matrix_view centroids, - MappingOpT mapping_op = raft::identity_op(), + MappingOpT mapping_op = raft::identity_op(), std::optional> inertia = std::nullopt) { RAFT_EXPECTS(X.extent(1) == centroids.extent(1), From 0a7f0261d3780c08f8ca65831d68ddcb52bc230c Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 9 Mar 2026 14:35:39 -0700 Subject: [PATCH 7/9] rm null dataset norm --- cpp/src/cluster/detail/kmeans_balanced.cuh | 11 +++++------ cpp/src/cluster/kmeans_balanced.cuh | 1 - 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index cab01221e7..5f2164f71a 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1012,7 +1012,6 @@ auto build_fine_clusters(const raft::resources& handle, * @param[out] cluster_centers A device pointer to the found cluster centers [n_clusters, dim] * @param[in] n_clusters Requested number of clusters * @param[in] mapping_op Mapping operation from T to MathT - * @param[in] dataset_norm (optional) Pre-computed L2 norms of each row in the dataset [n_rows] * @param[out] inertia (optional) If non-null, the sum of squared distances of samples to * their closest cluster center is written here. * Only supported when T == MathT (float/double). @@ -1026,8 +1025,7 @@ void build_hierarchical(const raft::resources& handle, MathT* cluster_centers, IdxT n_clusters, MappingOpT mapping_op, - const MathT* dataset_norm = nullptr, - MathT* inertia = nullptr) + MathT* inertia = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); using LabelT = uint32_t; @@ -1046,9 +1044,10 @@ void build_hierarchical(const raft::resources& handle, // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); - if (dataset_norm == nullptr && (params.metric == cuvs::distance::DistanceType::L2Expanded || - params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || - params.metric == cuvs::distance::DistanceType::CosineExpanded)) { + const MathT* dataset_norm = nullptr; + if ((params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded)) { dataset_norm_buf.resize(n_rows, stream); for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index 2b250b92cf..0c0df03397 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -91,7 +91,6 @@ void fit(const raft::resources& handle, centroids.data_handle(), centroids.extent(0), mapping_op, - static_cast(nullptr), inertia_ptr); } From 10be5c41650bdd64eceff14b47b3b9426c9cc0f4 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 9 Mar 2026 14:40:45 -0700 Subject: [PATCH 8/9] add warning when T and MathT are different --- cpp/src/cluster/detail/kmeans_balanced.cuh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 5f2164f71a..f5dc759725 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1175,6 +1175,8 @@ void build_hierarchical(const raft::resources& handle, raft::make_device_matrix_view(cluster_centers, n_clusters, dim); cuvs::cluster::kmeans::cluster_cost( handle, X_view, centroids_view, raft::make_host_scalar_view(inertia)); + } else { + RAFT_LOG_WARN("Inertia is not computed for non float/double types"); } } } From 6a2a681c3fc48c170102be776f4c5f9ca3b35220 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 9 Mar 2026 14:54:12 -0700 Subject: [PATCH 9/9] use raft::mul_op --- cpp/src/cluster/kmeans.cuh | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 8e26372b46..e4f9821990 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -455,12 +455,11 @@ void cluster_cost( // Apply sample weights if provided if (sample_weight.has_value()) { - raft::linalg::map( - handle, - min_cluster_distance.view(), - [] __device__(DataT d, DataT w) { return d * w; }, - raft::make_const_mdspan(min_cluster_distance.view()), - sample_weight.value()); + raft::linalg::map(handle, + min_cluster_distance.view(), + raft::mul_op{}, + raft::make_const_mdspan(min_cluster_distance.view()), + sample_weight.value()); } auto device_cost = raft::make_device_scalar(handle, DataT(0));