Enables instantiating entropy model in graph mode with compression=True.

PiperOrigin-RevId: 314736009
Change-Id: I02322aa86358d428eb7155ad7f3818c53e1a1f17

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -54,10 +54,8 @@ def __init__(self, prior, coding_rank, compression=False,
         unit. Each coding unit is compressed to its own bit string, and the
         `bits()` method sums over each coding unit.
       compression: Boolean. If set to `True`, the range coding tables used by
-        `compress()` and `decompress()` will be built on instantiation. This
-        assumes eager mode (throws an error if in graph mode or inside a
-        `tf.function` call). If set to `False`, these two methods will not be
-        accessible.
+        `compress()` and `decompress()` will be built on instantiation. If set
+        to `False`, these two methods will not be accessible.
       likelihood_bound: Float. Lower bound for likelihood values, to prevent
         training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
@@ -81,8 +79,6 @@ def __init__(self, prior, coding_rank, compression=False,
       self._tail_mass = float(tail_mass)
       self._range_coder_precision = int(range_coder_precision)
       if self.compression:
-        if not tf.executing_eagerly():
-          raise RuntimeError("`compression=True` requires eager execution.")
         self._build_tables(prior)
 
   @property
@@ -207,7 +203,7 @@ def _build_tables(self, prior):
     # Sample the densities in the computed ranges, possibly computing more
     # samples than necessary at the upper end.
     max_length = tf.math.reduce_max(pmf_length)
-    if max_length > 2048:
+    if tf.executing_eagerly() and max_length > 2048:
       logging.warning(
           "Very wide PMF with %d elements may lead to out of memory issues. "
           "Consider priors with smaller dispersion or increasing `tail_mass` "

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -92,10 +92,8 @@ def __init__(self, prior, coding_rank, compression=False,
         unit. Each coding unit is compressed to its own bit string, and the
         `bits()` method sums over each coding unit.
       compression: Boolean. If set to `True`, the range coding tables used by
-        `compress()` and `decompress()` will be built on instantiation. This
-        assumes eager mode (throws an error if in graph mode or inside a
-        `tf.function` call). If set to `False`, these two methods will not be
-        accessible.
+        `compress()` and `decompress()` will be built on instantiation. If set
+        to `False`, these two methods will not be accessible.
       likelihood_bound: Float. Lower bound for likelihood values, to prevent
         training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
@@ -119,7 +117,10 @@ def __init__(self, prior, coding_rank, compression=False,
       # Optimization: if the quantization offset is zero, we don't need to
       # subtract/add it when quantizing, and we don't need to serialize its
       # value. Note that this code will only work in eager mode.
-      if tf.reduce_all(tf.equal(quantization_offset, 0.)):
+      # TODO(jonycgn): Reconsider if this optimization is worth keeping once
+      # the implementation is stable.
+      if tf.executing_eagerly() and tf.reduce_all(
+          tf.equal(quantization_offset, 0.)):
         quantization_offset = None
       else:
         quantization_offset = tf.broadcast_to(
@@ -260,6 +261,8 @@ def decompress(self, strings, broadcast_shape):
       A `tf.Tensor` of shape `strings.shape + broadcast_shape +
       self.prior_shape`.
     """
+    strings = tf.convert_to_tensor(strings, dtype=tf.string)
+    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)

tensorflow_compression/python/entropy_models/continuous_batched_test.py

@@ -154,5 +154,28 @@ def test_compression_works_after_serialization_no_offset(self):
     self.assertAllEqual(em.compress(x), x_compressed)
     self.assertAllEqual(em.decompress(x_compressed, [100]), x_quantized)
 
+  def test_compression_works_in_tf_function(self):
+    noisy = uniform_noise.NoisyNormal(loc=0, scale=5.)
+    sample = noisy.base.sample([100])
+
+    # Since tf.function traces each function twice, and only allows variable
+    # creation in the first call, we need to have a stateful object in which we
+    # create the entropy model only the first time the function is called, and
+    # store it for the second time.
+
+    class Compressor(object):
+
+      def compress(self, values):
+        if not hasattr(self, "em"):
+          self.em = ContinuousBatchedEntropyModel(noisy, 1, compression=True)
+        compressed = self.em.compress(values)
+        decompressed = self.em.decompress(compressed, [])
+        return decompressed
+
+    values_eager = Compressor().compress(sample)
+    values_function = tf.function(Compressor().compress)(sample)
+    self.assertAllEqual(values_eager, values_function)
+
+
 if __name__ == "__main__":
   tf.test.main()