Lower bounds the overflow bin by zero probability.

PiperOrigin-RevId: 323625009
Change-Id: I44f7c2f6120e492978eda252d5f981b21edf8a2d

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -237,7 +237,8 @@ def _build_tables(self, prior):
       def loop_body(args):
         prob, length = args
         prob = prob[:length]
-        prob = tf.concat([prob, 1 - tf.reduce_sum(prob, keepdims=True)], axis=0)
+        overflow = tf.math.maximum(1 - tf.reduce_sum(prob, keepdims=True), 0.)
+        prob = tf.concat([prob, overflow], axis=0)
         cdf = range_coding_ops.pmf_to_quantized_cdf(
             prob, precision=self.range_coder_precision)
         return tf.pad(