limit iterations

Author: tarang-jain

cpp/src/cluster/detail/kmeans_common.cuh

@@ -548,10 +548,8 @@ void compute_centroid_adjustments(
     // Materialize the (possibly lazy) labels iterator into a device buffer so
     // we can read it on host without disturbing the original sequence.
     rmm::device_uvector<IndexT> mat_labels(static_cast<std::size_t>(n_samples), stream);
-    thrust::copy_n(raft::resource::get_thrust_policy(handle),
-                   cluster_labels,
-                   n_samples,
-                   mat_labels.data());
+    thrust::copy_n(
+      raft::resource::get_thrust_policy(handle), cluster_labels, n_samples, mat_labels.data());
 
     const std::size_t n_X_elems =
       static_cast<std::size_t>(n_samples) * static_cast<std::size_t>(n_features);

cpp/src/cluster/detail/minClusterDistanceCompute.cu

@@ -154,9 +154,7 @@ void minClusterAndDistanceCompute(
   // whether outputs are bitwise identical across the K runs.
   {
     static std::atomic<bool> s_self_test_done{false};
-    if (!s_self_test_done.exchange(true)) {
-      runFusedNNDeterminismSelfTest<DataT, IndexT>(handle);
-    }
+    if (!s_self_test_done.exchange(true)) { runFusedNNDeterminismSelfTest<DataT, IndexT>(handle); }
   }
 
   // ----- Nondeterminism diagnostic: fingerprint INPUTS (X, centroids, L2NormX) -----

cpp/tests/cluster/kmeans.cu

@@ -381,6 +381,9 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam<KmeansBatchedInputs
     params.tol                 = testparams.tol;
     params.rng_state.seed      = 1;
     params.oversampling_factor = 0;
+    // Limit the number of iterations to ensure same number of iterations for reference and batched
+    // code paths.
+    params.max_iter = 3;
 
     auto stream = raft::resource::get_cuda_stream(handle);
 
@@ -656,7 +659,7 @@ const std::vector<KmeansBatchedInputs<float>> batched_inputsf2 = {
   {1000, 64, 5, 0.0001f, false, 500},
   {1000, 100, 20, 0.0001f, true, 30},
   {1000, 10, 20, 0.0001f, false, 30},
-  {10000, 16, 10, 0.001f, true, 1000},
+  {10000, 16, 10, 0.0001f, true, 1000},
   {10000, 96, 10, 0.0001f, false, 10000},
 };