tf.summary.histogram allocates a lot of memory when called with large data

Environment information (required)

Irrelevant

Issue description

I am calling tf.summary.histogram with large enough data (in my case a 1-D tensor with ~16M int32 elements). I noticed that this would cause the method to allocate around 3.6GB of memory, with the following log line from tensorflow:

2024-01-02 19:27:52.663452: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 3903564240 exceeds 10% of free system memory.

Digging down deeper, I notice that this call to tf.one_hot in the method to calculate buckets is the culprit. This tensor has bucket_count times the original number of elements and with it being a tf.float64 type, it explains the 3.6GB of memory allocation in my case with default bucket count value of 30 (16M * 30 * 8 = ~3.6GB).

Looks like this tensor is only used to calculate bucket_counts tensor, which is a count of the number of elements in each bucket. Can this be done in a more memory efficient way? Something like:

bucket_counts = tf.map_fn(lambda i: tf.reduce_sum(tf.cast(clamped_indices == i, tf.int32)), tf.range(bucket_count))

Start in tensorboard/plugins/histogram/summary_v2.py at tf.summary.histogram and the bucket calculation around the tf.one_hot call. Check how the intermediate tensor produces bucket_counts for a 1-D tensor with about 16M int32 elements. Done means preserving the histogram bucket counts without the reported multi-gigabyte allocation, with the relevant histogram behavior verified.