Discrete histograms/counters across steps/epochs

I'm currently tracking several discrete values, e.g. the number of iterations a loop has gone through, over several training steps and dumping them to a tf.summary.histogram at the end of each epoch.

My current approach amounts to storing counts of the values seen in a counts = tf.Variable(shape=[max_counter_value], initializer='zeros', dtype=tf.int32, trainable=False) and then constructing a "fake" tf.summary_histogram, e.g.

      data = tf.repeat(
          tf.range(max_counter_value, dtype=tf.float32),
          repeats=counts)
      tf.summary.histogram(name, data, step=step, buckets=2 * max_counter_value - 1)

This is a bit clunky for several reasons:

• I need to inflate the counts into the data Tensor, which can contain several thousands/millions of values, even though tf.summary.histogram probably re-computes counts internally,
• Plotting the data in TensorBoard smears the discrete values over several buckets, so I never actually see the total count,
• (Somewhat tangentially) I have to handle accumulating and flushing the counters myself.

What would be really super cool would be a function named something like tf.summary.counter(value: tf.Tensor, name: str, flush_every_n_epochs:int = 1) where I can just dump in Tensors of integer types and get the discrete (unsmoothed) histograms every 'n' epochs.

I'm guessing the third part (accumulating values across steps) is probably a bit iffy since it would require maintaining some kind of state, but I'm hoping something like calling tf.summary.histogram with a set of pre-bucketed counts should be possible?

Cheers, Pedro