Allow disabling sanity check.

This makes it possible to feed random bits into the decoder
to treat it like a generator and produce samples from the
learned data model.

PiperOrigin-RevId: 438913987
Change-Id: Ia46879dff277b82cafaa806555658edf89a4fe84

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -122,6 +122,7 @@ def __init__(self,
                cdf_shapes=None,
                offset_heuristic=True,
                quantization_offset=None,
+               decode_sanity_check=True,
                laplace_tail_mass=0):
     """Initializes the instance.
 
@@ -168,6 +169,8 @@ def __init__(self,
         if you are using soft quantization during training.
       quantization_offset: `tf.Tensor` or `None`. The quantization offsets to
         use. If provided (not `None`), then `offset_heuristic` is ineffective.
+      decode_sanity_check: Boolean. If `True`, an raises an error if the binary
+        strings passed into `decompress` are not completely decoded.
       laplace_tail_mass: Float. If positive, will augment the prior with a
         Laplace mixture for training stability. (experimental)
     """
@@ -197,6 +200,7 @@ def __init__(self,
         prior_shape if prior is None else prior.batch_shape)
     if self.coding_rank < self.prior_shape.rank:
       raise ValueError("`coding_rank` can't be smaller than `prior_shape`.")
+    self.decode_sanity_check = decode_sanity_check
 
     with self.name_scope:
       if cdf_shapes is not None:
@@ -402,7 +406,8 @@ def decompress(self, strings, broadcast_shape):
     handle, symbols = gen_ops.entropy_decode_channel(
         handle, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
-    tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+    if self.decode_sanity_check:
+      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
     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

@@ -138,6 +138,7 @@ def __init__(self,
                range_coder_precision=12,
                bottleneck_dtype=None,
                prior_dtype=tf.float32,
+               decode_sanity_check=True,
                laplace_tail_mass=0):
     """Initializes the instance.
 
@@ -186,6 +187,8 @@ def __init__(self,
         Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
       prior_dtype: `tf.dtypes.DType`. Data type of prior and probability
         computations. Defaults to `tf.float32`.
+      decode_sanity_check: Boolean. If `True`, an raises an error if the binary
+        strings passed into `decompress` are not completely decoded.
       laplace_tail_mass: Float. If positive, will augment the prior with a
         laplace mixture for training stability. (experimental)
     """
@@ -216,6 +219,7 @@ def __init__(self,
     self._prior_fn = prior_fn
     self._parameter_fns = dict(parameter_fns)
     self._prior_dtype = tf.as_dtype(prior_dtype)
+    self.decode_sanity_check = decode_sanity_check
 
     with self.name_scope:
       if self.compression:
@@ -404,7 +408,8 @@ def decompress(self, strings, indexes):
     handle, symbols = gen_ops.entropy_decode_index(
         handle, flat_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
-    tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+    if self.decode_sanity_check:
+      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
     symbols += tf.gather(self.cdf_offset, flat_indexes)
     return tf.cast(symbols, self.bottleneck_dtype)
 

tensorflow_compression/python/entropy_models/universal.py

@@ -82,7 +82,8 @@ def __init__(self,
                range_coder_precision=12,
                bottleneck_dtype=None,
                num_noise_levels=15,
-               stateless=False):
+               stateless=False,
+               decode_sanity_check=True):
     """Initializes the instance.
 
     Args:
@@ -118,6 +119,8 @@ def __init__(self,
         allows it to be constructed within a `tf.function` body. If
         `compression=False`, then `stateless=True` is implied and the provided
         value is ignored.
+      decode_sanity_check: Boolean. If `True`, an raises an error if the binary
+        strings passed into `decompress` are not completely decoded.
     """
     if prior.event_shape.rank:
       raise ValueError("`prior` must be a (batch of) scalar distribution(s).")
@@ -135,6 +138,7 @@ def __init__(self,
     self._num_noise_levels = num_noise_levels
     if self.coding_rank < self.prior_shape.rank:
       raise ValueError("`coding_rank` can't be smaller than `prior_shape`.")
+    self.decode_sanity_check = decode_sanity_check
 
     with self.name_scope:
       if self.compression:
@@ -285,7 +289,8 @@ def decompress(self, strings, broadcast_shape):
     handle, symbols = gen_ops.entropy_decode_index(
         handle, decode_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
-    tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+    if self.decode_sanity_check:
+      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
     symbols += tf.gather(self.cdf_offset, indexes)
     outputs = tf.cast(symbols, self.bottleneck_dtype)
     return outputs + offset
@@ -321,7 +326,8 @@ def __init__(self,
                bottleneck_dtype=None,
                prior_dtype=tf.float32,
                stateless=False,
-               num_noise_levels=15):
+               num_noise_levels=15,
+               decode_sanity_check=True):
     """Initializes the instance.
 
     Args:
@@ -364,6 +370,8 @@ def __init__(self,
         rather than `Variable`s.
       num_noise_levels: Integer. The number of levels used to quantize the
         uniform noise.
+      decode_sanity_check: Boolean. If `True`, an raises an error if the binary
+        strings passed into `decompress` are not completely decoded.
     """
     if coding_rank <= 0:
       raise ValueError("`coding_rank` must be larger than 0.")
@@ -393,6 +401,7 @@ def __init__(self,
     self._parameter_fns = dict(parameter_fns)
     self._prior_dtype = tf.as_dtype(prior_dtype)
     self._num_noise_levels = num_noise_levels
+    self.decode_sanity_check = decode_sanity_check
 
     with self.name_scope:
       if self.compression:
@@ -577,7 +586,8 @@ def decompress(self, strings, indexes):
     handle, symbols = gen_ops.entropy_decode_index(
         handle, flat_indexes, decode_shape, self.cdf_offset.dtype)
     sanity = gen_ops.entropy_decode_finalize(handle)
-    tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
+    if self.decode_sanity_check:
+      tf.debugging.assert_equal(sanity, True, message="Sanity check failed.")
     symbols += tf.gather(self.cdf_offset, flat_indexes)
     offset = self._offset_from_indexes(indexes)
     return tf.cast(symbols, self.bottleneck_dtype) + offset