histogram: implement buckets_v3

[yatbear]

Main changes:

• For single value input, the histogram output will no longer be reshaped to length 1, instead, there will be {bucket_count} buckets where the first {bucket_count - 1} buckets have zero counts and the last bucket has count {data_size}, since only the last bucket is closed-closed and the other buckets are closed-open.
• If bucket_count is invalid (<= 0), an empty tensor of shape (0, 3) will be returned.
• Naming singular is replaced by single_value in v3 to avoid confusion with the mathematical term singularity.

Tests for v3 pb op will be added once histogram_pb_v3 is implemented.

#histogram #tpu

[nfelt]

So I think the reason this was failing is that in graph mode, when it compiles the TensorFlow code into a graph, all of the branches of the tf.cond()s are inserted as parts of the graph (even if those branches won't be taken when the graph is actually executed). When doing that, TensorFlow does shape inference, described a little bit here: https://www.tensorflow.org/guide/create_op#shape_functions_in_c

Basically, it tries to compute what the resulting shape will be for all of the ops in all of the branches. At the line where we currently do tf.fill([bucket_count - 1], 0), because bucket_count is a constant, it is able to determine at compile time that this will result in tf.fill([-1], 0) and the shape function for tf.fill emits an error, because this would result in an output shape that doesn't make sense.

Hence the error (at least, the one that I was seeing when trying this PR locally):

ValueError: Fill dimensions must be >= 0 for '{{node zero_bucket_count/write_summary/buckets/cond/cond/Fill_1}} = Fill[T=DT_INT32, index_type=DT_INT32](zero_bucket_count/write_summary/buckets/cond/cond/Fill_1/dims, zero_bucket_count/write_summary/buckets/cond/cond/Fill_1/value)' with input shapes: [1], [] and with input tensors computed as partial shapes: input[0] = [?].

In other words, it doesn't actually mean it's picking the wrong branch of the conditional to execute; it's that during graph compilation, it visits all the branches regardless of what will be ultimately executed.

(Aside: you could argue that when bucket_count is a constant 0, we could actually optimize away all the tf.cond ops and all the branches besides the empty case, at which point we wouldn't hit this error. I'm guessing that the compilation's constant-folding isn't quite smart enough to do that. Also, if you change the test to pass bucket_count=tf.Variable(0) and replace the use of max() with tf.math.maximum() as mentioned in the other comment, it will also pass, since now that the bucket count is a variable, it doesn't do as much shape inference and can't raise the false positive error.)

Ultimately, the best way to fix this is to just ensure that all the branches are still creating the correctly shaped tensor of dimensions (bucket_count, 3), even when we know that branch will never be executed. Here's a patch that gets the test to pass:

diff --git i/tensorboard/plugins/histogram/summary_v2.py w/tensorboard/plugins/histogram/summary_v2.py
index bff980492..e1cd01846 100644
--- i/tensorboard/plugins/histogram/summary_v2.py
+++ w/tensorboard/plugins/histogram/summary_v2.py
@@ -425,6 +425,8 @@ def _buckets_v3(data, bucket_count=None):
     with tf.name_scope("buckets"):
         tf.debugging.assert_scalar(bucket_count)
         tf.debugging.assert_type(bucket_count, tf.int32)
+        # Treat a negative bucket count as zero.
+        bucket_count = tf.math.maximum(0, bucket_count)
         data = tf.reshape(data, shape=[-1])  # flatten
         data = tf.cast(data, tf.float64)
         data_size = tf.size(input=data)
@@ -440,7 +442,7 @@ def _buckets_v3(data, bucket_count=None):
             2. If the input data is empty, a tensor of shape (bucket_count, 3)
               of all zero values will be returned.
             """
-            return tf.zeros((max(0, bucket_count), 3), dtype=tf.float64)
+            return tf.zeros((bucket_count, 3), dtype=tf.float64)

         def when_nonempty():
             min_ = tf.reduce_min(input_tensor=data)
@@ -480,11 +482,12 @@ def _buckets_v3(data, bucket_count=None):
                 """When input data contains a single unique value."""
                 # Left and right edges are the same for single value input.
                 edges = tf.fill([bucket_count], max_)
-                # Counts for the first {bucket_count - 1} buckets [v, v) are 0.
-                zero_bucket_counts = tf.fill([bucket_count - 1], 0)
-                # Count for last bucket [v, v] is {data_size}.
+                # Bucket counts are 0 except the last bucket (if bucket_count > 0),
+                # which is `data_size`. Ensure that the resulting counts vector has
+                # length `bucket_count` always, including the bucket_count==0 case.
+                zeroes = tf.fill([bucket_count], 0)
                 bucket_counts = tf.cast(
-                    tf.concat([zero_bucket_counts, [data_size]], 0),
+                    tf.concat([zeroes[:-1], [data_size]], 0)[:bucket_count],
                     dtype=tf.float64,
                 )
                 return tf.transpose(a=tf.stack([edges, edges, bucket_counts]))

[yatbear]

Thanks so much for the clear explanation and the patch!

[nfelt]

We should use tf.math.maximum() here, rather than Python's max(), in order to be compatible with bucket_count values that aren't native Python numbers (e.g. a tf.constant() instead).

Also I'd recommend just adding a conversion up at the top that does bucket_count = tf.math.maximum(0, bucket_count) so we don't have to worry about negative bucket counts in any of the later logic.

[yatbear]

Done. Thanks!

[nfelt]

LGTM modulo what looks like unintentional additions

[nfelt]

Unintentional? Also the autograph=True addition above seems like it shouldn't be necessary? (I believe it's true by default.)

[yatbear]

Removed. Thanks!