Merge pull request #128741 from destro4evr-rgb:fix/kmeans-plusplus-oob-read

PiperOrigin-RevId: 1003694965

tensorflow/core/kernels/clustering_ops.cc

@@ -108,11 +108,6 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
     const int64_t point_dimensions = points_tensor.dim_size(1);
     const int64_t num_to_sample = num_to_sample_tensor.scalar<int64_t>()();
     const int64_t seed = seed_tensor.scalar<int64_t>()();
-    const int64_t num_retries_per_sample = [&]() {
-      const int64_t value = num_retries_per_sample_tensor.scalar<int64_t>()();
-      return value >= 0 ? value
-                        : 2 + static_cast<int64_t>(std::log(num_to_sample));
-    }();
 
     OP_REQUIRES(context, num_points > 0,
                 absl::InvalidArgumentError("Expected points.rows() > 0."));
@@ -123,6 +118,12 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
                     absl::StrCat("Expected num_to_sample <= points.rows(). ",
                                  num_to_sample, " vs ", num_points, ".")));
 
+    const int64_t num_retries_per_sample = [&]() {
+      const int64_t value = num_retries_per_sample_tensor.scalar<int64_t>()();
+      return value >= 0 ? value
+                        : 2 + static_cast<int64_t>(std::log(num_to_sample));
+    }();
+
     Tensor* output_sampled_points_tensor;
     OP_REQUIRES_OK(context,
                    context->allocate_output(
@@ -156,16 +157,27 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
 
     auto draw_one_sample = [&]() -> int64_t {
       if (sampled_indices.empty()) return rng.Uniform64(num_points);
+      const float total = min_distances_cumsum(num_points - 1);
       int64_t index = 0;
       do {
-        // If v is drawn from Uniform[0, distances.sum()), then
-        // Prob[cumsum(distances)(i - 1) <= v < cumsum(distances)(i)] is
-        // proportional to distances(i).
-        index = std::upper_bound(
-                    min_distances_cumsum.data(),
-                    min_distances_cumsum.data() + num_points,
-                    rng.RandFloat() * min_distances_cumsum(num_points - 1)) -
-                min_distances_cumsum.data();
+        if (!(total > 0.0f) || std::isinf(total)) {
+          // total is zero, negative, NaN, or infinity (e.g. all points
+          // identical, NaN distances, or overflow). Fall back to uniform
+          // sampling to avoid upper_bound returning end-of-range and an
+          // out-of-bounds index.
+          index = rng.Uniform64(num_points);
+        } else {
+          // If v is drawn from Uniform[0, distances.sum()), then
+          // Prob[cumsum(distances)(i - 1) <= v < cumsum(distances)(i)] is
+          // proportional to distances(i).
+          index = std::upper_bound(min_distances_cumsum.data(),
+                                   min_distances_cumsum.data() + num_points,
+                                   rng.RandFloat() * total) -
+                  min_distances_cumsum.data();
+          // upper_bound guarantees index < num_points when total > 0, but
+          // clamp defensively against floating-point rounding edge cases.
+          index = std::min(index, num_points - 1);
+        }
       } while (sampled_indices.find(index) != sampled_indices.end());
       return index;
     };
@@ -175,6 +187,7 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
       min_distances = min_distances.cwiseMin(GetHalfSquaredDistancesToY(
           points, points_half_squared_norm, points.row(sampled_index),
           points_half_squared_norm(sampled_index)));
+      min_distances(sampled_index) = 0.0f;
       return sampled_index;
     };
 
@@ -182,14 +195,15 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
       Eigen::VectorXf best_new_min_distances(num_points);
       float best_potential = std::numeric_limits<float>::infinity();
       int64_t best_sampled_index = 0;
-      for (int i = 1 + num_retries_per_sample; i > 0; --i) {
+      for (int64_t i = 0; i <= num_retries_per_sample; ++i) {
         const int64_t sampled_index = draw_one_sample();
         Eigen::VectorXf new_min_distances =
             min_distances.cwiseMin(GetHalfSquaredDistancesToY(
                 points, points_half_squared_norm, points.row(sampled_index),
                 points_half_squared_norm(sampled_index)));
+        new_min_distances(sampled_index) = 0.0f;
         const float potential = new_min_distances.sum();
-        if (potential < best_potential) {
+        if (i == 0 || potential < best_potential) {
           best_potential = potential;
           best_sampled_index = sampled_index;
           best_new_min_distances.swap(new_min_distances);
@@ -201,6 +215,7 @@ class KmeansPlusPlusInitializationOp : public OpKernel {
 
     for (int64_t i = 0; i < num_to_sample; ++i) {
       if (i > 0) {
+        min_distances = min_distances.cwiseMax(0.0f);
         std::partial_sum(min_distances.data(),
                          min_distances.data() + num_points,
                          min_distances_cumsum.data());

tensorflow/python/ops/clustering_ops_test.py

@@ -52,6 +52,47 @@ class KmeansPlusPlusInitializationTest(test.TestCase):
     for seed in range(100):
       self.runTestWithSeed(seed)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testIdenticalPoints(self):
+    points = np.array([[1.0, 2.0], [1.0, 2.0]], dtype=np.float32)
+    with self.cached_session():
+      for retries in (0, 3):
+        sampled = clustering_ops.kmeans_plus_plus_initialization(
+            points, num_to_sample=2, seed=12345, num_retries_per_sample=retries
+        )
+        result = self.evaluate(sampled)
+        self.assertEqual(result.shape, (2, 2))
+        self.assertAllClose(result, points)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNaNPoints(self):
+    points = np.full((2, 2), float("nan"), dtype=np.float32)
+    with self.cached_session():
+      for retries in (0, 3):
+        sampled = clustering_ops.kmeans_plus_plus_initialization(
+            points, num_to_sample=2, seed=12345, num_retries_per_sample=retries
+        )
+        result = self.evaluate(sampled)
+        self.assertEqual(result.shape, (2, 2))
+        self.assertEqual(result.dtype, np.float32)
+        self.assertTrue(np.isnan(result).all())
+
+  @test_util.run_in_graph_and_eager_modes
+  def testOverflowPoints(self):
+    # Values large enough that squared distance overflows float32 to +inf
+    # without individual norms overflowing to +inf (which would yield NaN),
+    # exercising the std::isinf(total) fallback in draw_one_sample.
+    points = np.array([[1e19, 1e19], [-1e19, -1e19]], dtype=np.float32)
+    with self.cached_session():
+      for retries in (0, 3):
+        sampled = clustering_ops.kmeans_plus_plus_initialization(
+            points, num_to_sample=2, seed=12345, num_retries_per_sample=retries
+        )
+        result = self.evaluate(sampled)
+        self.assertEqual(result.shape, (2, 2))
+        self.assertEqual(result.dtype, np.float32)
+        self.assertCountEqual(result.tolist(), points.tolist())
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class KMC2InitializationTest(test.TestCase):