Cleans up a few instances of Tensor type casting.

Adds an attribute `prior_shape_tensor`, analogous to the shape attributes on tfp
Distributions.

PiperOrigin-RevId: 320462150
Change-Id: I5660cfcaf61db231054481baf14c5837903fc0c7

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -122,6 +122,11 @@ def prior_shape(self):
     """Batch shape of `prior` (dimensions which are not assumed i.i.d.)."""
     return self._prior_shape
 
+  @property
+  def prior_shape_tensor(self):
+    """Batch shape of `prior` as a `Tensor`."""
+    return tf.constant(self.prior_shape, dtype=tf.int32)
+
   @property
   def coding_rank(self):
     """Number of innermost dimensions considered a coding unit."""
@@ -217,10 +222,10 @@ def _build_tables(self, prior):
     pmf = tf.reshape(pmf, [max_length, -1])
     pmf = tf.transpose(pmf)
 
-    pmf_length = tf.broadcast_to(pmf_length, self.prior_shape)
+    pmf_length = tf.broadcast_to(pmf_length, self.prior_shape_tensor)
     pmf_length = tf.reshape(pmf_length, [-1])
     cdf_length = pmf_length + 2
-    cdf_offset = tf.broadcast_to(-minima, self.prior_shape)
+    cdf_offset = tf.broadcast_to(-minima, self.prior_shape_tensor)
     cdf_offset = tf.reshape(cdf_offset, [-1])
 
     # Prevent tensors from bouncing back and forth between host and GPU.

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -124,7 +124,7 @@ def __init__(self, prior, coding_rank, compression=False,
         quantization_offset = None
       else:
         quantization_offset = tf.broadcast_to(
-            quantization_offset, self.prior_shape)
+            quantization_offset, self.prior_shape_tensor)
         quantization_offset = tf.Variable(
             quantization_offset, trainable=False, name="quantization_offset")
     self._quantization_offset = quantization_offset
@@ -133,9 +133,9 @@ def _compute_indexes(self, broadcast_shape):
     # TODO(jonycgn, ssjhv): Investigate broadcasting in range coding op.
     prior_size = functools.reduce(lambda x, y: x * y, self.prior_shape, 1)
     indexes = tf.range(prior_size, dtype=tf.int32)
-    indexes = tf.reshape(indexes, self.prior_shape)
+    indexes = tf.reshape(indexes, self.prior_shape_tensor)
     indexes = tf.broadcast_to(
-        indexes, tf.concat([broadcast_shape, self.prior_shape], 0))
+        indexes, tf.concat([broadcast_shape, self.prior_shape_tensor], 0))
     return indexes
 
   @tf.Module.with_name_scope
@@ -164,7 +164,8 @@ def bits(self, bottleneck, training=True):
     probs = self.prior.prob(quantized)
     probs = math_ops.lower_bound(probs, self.likelihood_bound)
     axes = tuple(range(-self.coding_rank, 0))
-    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / -tf.math.log(2.)
+    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / (
+        -tf.math.log(tf.constant(2., dtype=probs.dtype)))
     return bits
 
   @tf.Module.with_name_scope
@@ -265,7 +266,7 @@ def decompress(self, strings, broadcast_shape):
     broadcast_shape = tf.convert_to_tensor(broadcast_shape, dtype=tf.int32)
     batch_shape = tf.shape(strings)
     symbols_shape = tf.concat(
-        [batch_shape, broadcast_shape, self.prior_shape], 0)
+        [batch_shape, broadcast_shape, self.prior_shape_tensor], 0)
 
     indexes = self._compute_indexes(broadcast_shape)
     strings = tf.reshape(strings, [-1])
@@ -321,7 +322,7 @@ def from_config(cls, config):
     with self.name_scope:
       # pylint:disable=protected-access
       if config["quantization_offset"]:
-        zeros = tf.zeros(self.prior_shape, dtype=self.dtype)
+        zeros = tf.zeros(self.prior_shape_tensor, dtype=self.dtype)
         self._quantization_offset = tf.Variable(
             zeros, name="quantization_offset")
       else:

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -282,7 +282,8 @@ def bits(self, bottleneck, indexes, training=True):
     probs = prior.prob(quantized)
     probs = math_ops.lower_bound(probs, self.likelihood_bound)
     axes = tuple(range(-self.coding_rank, 0))
-    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / -tf.math.log(2.)
+    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / (
+        -tf.math.log(tf.constant(2., dtype=probs.dtype)))
     return bits
 
   @tf.Module.with_name_scope