histogram: cast tf.reduce_sum input to float64 (tensorflow#5337)

Author: yatbear

tensorboard/plugins/histogram/summary.py

@@ -81,7 +81,12 @@ def when_nonsingular():
                     tf.floor(offsets / bucket_width), dtype=tf.int32
                 )
                 clamped_indices = tf.minimum(bucket_indices, bucket_count - 1)
-                one_hots = tf.one_hot(clamped_indices, depth=bucket_count)
+                # Use float64 instead of float32 to avoid accumulating floating point error
+                # later in tf.reduce_sum when summing more than 2^24 individual `1.0` values.
+                # See https://github.com/tensorflow/tensorflow/issues/51419 for details.
+                one_hots = tf.one_hot(
+                    clamped_indices, depth=bucket_count, dtype=tf.float64
+                )
                 bucket_counts = tf.cast(
                     tf.reduce_sum(input_tensor=one_hots, axis=0),
                     dtype=tf.float64,

tensorboard/plugins/histogram/summary_test.py

@@ -119,6 +119,16 @@ def test_when_bucket_count_not_statically_known(self):
         buckets = tensor_util.make_ndarray(pb.value[0].tensor)
         self.assertEqual(buckets.shape, (bucket_count, 3))
 
+    def test_with_large_counts(self):
+        # Check for accumulating floating point errors with large counts (> 2^24).
+        # See https://github.com/tensorflow/tensorflow/issues/51419 for details.
+        large_count = 20_000_000
+        data = [0] + [1] * large_count
+        pb = self.histogram("large_count", data=data, buckets=2)
+        buckets = tensor_util.make_ndarray(pb.value[0].tensor)
+        self.assertEqual(buckets[0][2], 1)
+        self.assertEqual(buckets[1][2], large_count)
+
 
 class SummaryV1PbTest(SummaryBaseTest, tf.test.TestCase):
     def histogram(self, *args, **kwargs):

tensorboard/plugins/histogram/summary_v2.py

@@ -214,7 +214,12 @@ def when_nonsingular():
                     tf.floor(offsets / bucket_width), dtype=tf.int32
                 )
                 clamped_indices = tf.minimum(bucket_indices, bucket_count - 1)
-                one_hots = tf.one_hot(clamped_indices, depth=bucket_count)
+                # Use float64 instead of float32 to avoid accumulating floating point error
+                # later in tf.reduce_sum when summing more than 2^24 individual `1.0` values.
+                # See https://github.com/tensorflow/tensorflow/issues/51419 for details.
+                one_hots = tf.one_hot(
+                    clamped_indices, depth=bucket_count, dtype=tf.float64
+                )
                 bucket_counts = tf.cast(
                     tf.reduce_sum(input_tensor=one_hots, axis=0),
                     dtype=tf.float64,