Small edits to distributions/entropy models.

- Lower bounds the probabilities computed in DeepFactorized by zero.
- Adds a name scope to the variables created in DeepFactorized.
- Enables deleting the prior distribution from the entropy model.
- Removes deprecated backprop argument from tf.while.

PiperOrigin-RevId: 323400592
Change-Id: Iff3cb1cb3f83fb7381798e6f71779126ac7a00da

Author: Johannes Ballé

tensorflow_compression/python/distributions/BUILD

@@ -19,7 +19,10 @@ py_library(
     name = "deep_factorized",
     srcs = ["deep_factorized.py"],
     srcs_version = "PY3",
-    deps = [":helpers"],
+    deps = [
+        ":helpers",
+        "//tensorflow_compression/python/ops:math_ops",
+    ],
 )
 
 py_test(

tensorflow_compression/python/distributions/deep_factorized.py

@@ -19,6 +19,7 @@
 import tensorflow_probability as tfp
 
 from tensorflow_compression.python.distributions import helpers
+from tensorflow_compression.python.ops import math_ops
 
 
 __all__ = ["DeepFactorized"]
@@ -77,7 +78,8 @@ def __init__(self, batch_shape=(), num_filters=(3, 3), init_scale=10,
         parameters=parameters,
         name=name,
     )
-    self._make_variables()
+    with self.name_scope:
+      self._make_variables()
 
   @property
   def num_filters(self):
@@ -168,6 +170,7 @@ def _prob(self, y):
     # Flip signs if we can move more towards the left tail of the sigmoid.
     sign = tf.stop_gradient(-tf.math.sign(lower + upper))
     p = abs(tf.sigmoid(sign * upper) - tf.sigmoid(sign * lower))
+    p = math_ops.lower_bound(p, 0.)
 
     # Convert back to (broadcasted) input tensor shape.
     p = tf.transpose(p, (2, 1, 0))

tensorflow_compression/python/distributions/helpers.py

@@ -82,8 +82,7 @@ def loop_body(tails, m, v, count):
     init_v = tf.ones(shape, dtype=dtype)
     init_count = tf.zeros(shape, dtype=tf.int32)
     return tf.while_loop(
-        loop_cond, loop_body, (init_tails, init_m, init_v, init_count),
-        back_prop=False)[0]
+        loop_cond, loop_body, (init_tails, init_m, init_v, init_count))[0]
 
 
 def quantization_offset(distribution):

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -88,9 +88,13 @@ def prior(self):
       raise RuntimeError(
           "This entropy model doesn't hold a reference to its prior "
           "distribution. This can happen when it is unserialized, because "
-          "the prior is generally not serializable.")
+          "the prior is not generally serializable.")
     return self._prior
 
+  @prior.deleter
+  def prior(self):
+    del self._prior
+
   def _check_compression(self):
     if not self.compression:
       raise RuntimeError(