Added tf.control_dependencies() context managers around tf.debugging.assert* ops.

This is for better compatibility with TF1. However, there may be some side
effects not executed in TF1 still.

PiperOrigin-RevId: 485946240
Change-Id: Icb3a7f4cb5c9edc5de18b0751837f000c1f51287

Author: ssjhv

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -301,17 +301,16 @@ def _log_prob(self, prior, bottleneck_perturbed):
     laplace_tail_mass = self.laplace_tail_mass
 
     def mixture_log_prob_fn():
-      tf.debugging.assert_less(
-          laplace_tail_mass,
-          tf.constant(1.0, prior.dtype),
-          message="`laplace_tail_mass` must be less than 1.")
-      laplace_prior = uniform_noise.NoisyLaplace(
-          loc=tf.constant(0, dtype=prior.dtype),
-          scale=tf.constant(1, dtype=prior.dtype))
-      probs = prior.prob(bottleneck_perturbed)
-      probs = ((1 - laplace_tail_mass) * probs +
-               laplace_tail_mass *
-               laplace_prior.prob(bottleneck_perturbed))
+      with tf.control_dependencies([tf.debugging.assert_less(
+          laplace_tail_mass, tf.constant(1.0, prior.dtype),
+          message="`laplace_tail_mass` must be less than 1.")]):
+        laplace_prior = uniform_noise.NoisyLaplace(
+            loc=tf.constant(0, dtype=prior.dtype),
+            scale=tf.constant(1, dtype=prior.dtype))
+        probs = prior.prob(bottleneck_perturbed)
+        probs = ((1 - laplace_tail_mass) * probs +
+                 laplace_tail_mass *
+                 laplace_prior.prob(bottleneck_perturbed))
       probs_too_small = probs < 1e-10
       probs_bounded = tf.maximum(probs, 1e-10)
       return tf.where(

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -410,7 +410,9 @@ def decompress(self, strings, broadcast_shape):
         handle, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
     if self.decode_sanity_check:
-      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+      with tf.control_dependencies([tf.debugging.assert_equal(
+          sanity, True, message="Sanity check failed.")]):
+        symbols = tf.identity(symbols)
     symbols += self.cdf_offset
     symbols = tf.reshape(symbols, output_shape)
     outputs = tf.cast(symbols, self.bottleneck_dtype)

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -410,7 +410,9 @@ def decompress(self, strings, indexes):
         handle, flat_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
     if self.decode_sanity_check:
-      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+      with tf.control_dependencies([tf.debugging.assert_equal(
+          sanity, True, message="Sanity check failed.")]):
+        symbols = tf.identity(symbols)
     symbols += tf.gather(self.cdf_offset, flat_indexes)
     return tf.cast(symbols, self.bottleneck_dtype)
 

tensorflow_compression/python/entropy_models/universal.py

@@ -292,7 +292,9 @@ def decompress(self, strings, broadcast_shape):
         handle, decode_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
     if self.decode_sanity_check:
-      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+      with tf.control_dependencies([tf.debugging.assert_equal(
+          sanity, True, message="Sanity check failed.")]):
+        symbols = tf.identity(symbols)
     symbols += tf.gather(self.cdf_offset, indexes)
     outputs = tf.cast(symbols, self.bottleneck_dtype)
     return outputs + offset
@@ -589,7 +591,9 @@ def decompress(self, strings, indexes):
         handle, flat_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
     if self.decode_sanity_check:
-      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+      with tf.control_dependencies([tf.debugging.assert_equal(
+          sanity, True, message="Sanity check failed.")]):
+        symbols = tf.identity(symbols)
     symbols += tf.gather(self.cdf_offset, flat_indexes)
     offset = self._offset_from_indexes(indexes)
     return tf.cast(symbols, self.bottleneck_dtype) + offset