Update (experimental) laplace_tail_mass to use NoisyLaplace instead of Laplace.

This is because in some cases we assume that the distribution is a valid PMF when
evaluated at integer-spaced locations.

PiperOrigin-RevId: 460402710
Change-Id: Idac61b6865453d906ee93334437905c4eadb750f

Author: relational

tensorflow_compression/python/distributions/uniform_noise.py

@@ -24,6 +24,7 @@
     "NoisyMixtureSameFamily",
     "NoisyNormal",
     "NoisyLogistic",
+    "NoisyLaplace",
     "NoisyNormalMixture",
     "NoisyLogisticMixture",
 ]
@@ -265,6 +266,14 @@ def __init__(self, name="NoisyLogistic", **kwargs):
     super().__init__(tfp.distributions.Logistic(**kwargs), name=name)
 
 
+class NoisyLaplace(UniformNoiseAdapter):
+  """Laplacian distribution with additive i.i.d. uniform noise."""
+
+  def __init__(self, name="NoisyLaplace", **kwargs):
+    """Initializer, taking the same arguments as `tfpd.Laplace`."""
+    super().__init__(tfp.distributions.Laplace(**kwargs), name=name)
+
+
 class NoisyNormalMixture(NoisyMixtureSameFamily):
   """Mixture of normal distributions with additive i.i.d. uniform noise."""
 

tensorflow_compression/python/distributions/uniform_noise_test.py

@@ -53,6 +53,13 @@ def test_uniform_is_special_case(self):
     x = tf.linspace(mean - 1, mean + 1, 10)
     self.assertAllClose(dist.prob(x), [0, 0, 0, 1, 1, 1, 1, 0, 0, 0])
 
+  def test_is_pmf_at_integer_grid(self):
+    # When evaluated at integers, P_i = p(c+i) should be a valid PMF.
+    dist = self.dist_cls(loc=0.1, scale=0.3)
+    x = 0.05 + tf.range(-100, 100, dtype=tf.float32)
+    pmf_probs = dist.prob(x)
+    self.assertAllClose(tf.reduce_sum(pmf_probs), 1.0)
+
   def test_sampling_works(self):
     dist = self.dist_cls(loc=0, scale=[3, 5])
     sample = dist.sample((5, 4))
@@ -90,6 +97,11 @@ class NoisyLogisticTest(LocationScaleTest, tf.test.TestCase):
   dist_cls = uniform_noise.NoisyLogistic
 
 
+class NoisyLaplaceTest(LocationScaleTest, tf.test.TestCase):
+
+  dist_cls = uniform_noise.NoisyLaplace
+
+
 class MixtureTest:
   """Common tests for noisy mixture distributions."""
 

tensorflow_compression/python/entropy_models/BUILD

@@ -20,6 +20,7 @@ py_library(
     srcs = ["continuous_base.py"],
     deps = [
         "//tensorflow_compression/python/distributions:helpers",
+        "//tensorflow_compression/python/distributions:uniform_noise",
         "//tensorflow_compression/python/ops:gen_ops",
     ],
 )

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -17,8 +17,8 @@
 import abc
 from absl import logging
 import tensorflow as tf
-import tensorflow_probability as tfp
 from tensorflow_compression.python.distributions import helpers
+from tensorflow_compression.python.distributions import uniform_noise
 from tensorflow_compression.python.ops import gen_ops
 
 
@@ -69,7 +69,7 @@ def __init__(self,
       bottleneck_dtype: `tf.dtypes.DType`. Data type of bottleneck tensor.
         Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
       laplace_tail_mass: Float. If non-zero, will augment the prior with a
-        Laplace mixture for training stability. (experimental)
+        `NoisyLaplace` mixture component for training stability. (experimental)
     """
     super().__init__()
     self._prior = None  # This will be set by subclasses, if appropriate.
@@ -137,7 +137,7 @@ def expected_grads(self):
 
   @property
   def laplace_tail_mass(self):
-    """Whether to augment the prior with a Laplace mixture."""
+    """Whether to augment the prior with a `NoisyLaplace` mixture."""
     return self._laplace_tail_mass
 
   @property
@@ -298,7 +298,7 @@ def _log_prob(self, prior, bottleneck_perturbed):
     """Evaluates prior.log_prob(bottleneck + noise)."""
     bottleneck_perturbed = tf.cast(bottleneck_perturbed, prior.dtype)
     if self.laplace_tail_mass:
-      laplace_prior = tfp.distributions.Laplace(
+      laplace_prior = uniform_noise.NoisyLaplace(
           loc=tf.constant(0, dtype=prior.dtype),
           scale=tf.constant(1, dtype=prior.dtype))
       probs = prior.prob(bottleneck_perturbed)

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -171,8 +171,8 @@ def __init__(self,
         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)
+      laplace_tail_mass: Float. If non-zero, will augment the prior with a
+        `NoisyLaplace` mixture component for training stability. (experimental)
     """
     if (prior is None) == (prior_shape is None):
       raise ValueError("Either `prior` or `prior_shape` must be provided.")

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -189,8 +189,8 @@ def __init__(self,
         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)
+      laplace_tail_mass: Float. If non-zero, will augment the prior with a
+        `NoisyLaplace` mixture component for training stability. (experimental)
     """
     if not callable(prior_fn):
       raise TypeError("`prior_fn` must be a class or factory function.")
@@ -496,8 +496,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`.
-      laplace_tail_mass: Float. If positive, will augment the prior with a
-        laplace mixture for training stability. (experimental)
+      laplace_tail_mass: Float. If non-zero, will augment the prior with a
+        `NoisyLaplace` mixture component for training stability. (experimental)
     """
     num_scales = int(num_scales)
     super().__init__(

tensorflow_compression/python/entropy_models/universal_test.py

@@ -98,7 +98,7 @@ def test_laplace_tail_mass_for_large_inputs(self):
         coding_rank=1,
         compression=True,
         laplace_tail_mass=1e-3)
-    x = tf.convert_to_tensor([1e3, 1e4, 1e5, 1e6, 1e7, 1e8], tf.float32)
+    x = tf.convert_to_tensor([1e3, 1e4, 1e5, 1e6], tf.float32)
     _, bits = em(x[..., None])
     self.assertAllClose(bits, tf.abs(x) / tf.math.log(2.0), rtol=0.01)