Fix instability in mixture distributions if log_prob becomes -inf.

PiperOrigin-RevId: 516755689
Change-Id: I866976f926256fcef583313f734c055e60ae452d

Author: Fabian Mentzer

tensorflow_compression/python/distributions/uniform_noise.py

@@ -44,7 +44,11 @@ def _logsum_expbig_minus_expsmall(big, small):
     `tf.Tensor` containing the result.
   """
   with tf.name_scope("logsum_expbig_minus_expsmall"):
-    return tf.math.log1p(-tf.exp(small - big)) + big
+    # Have to special case `inf` and `-inf` since otherwise we get a NaN
+    # out of the exp (if both small and big are -inf).
+    return tf.where(
+        tf.math.is_inf(big), big, tf.math.log1p(-tf.exp(small - big)) + big
+    )
 
 
 class UniformNoiseAdapter(tfp.distributions.Distribution):

tensorflow_compression/python/distributions/uniform_noise_test.py

@@ -161,6 +161,15 @@ def test_stats_throw_error(self):
     with self.assertRaises(NotImplementedError):
       dist.survival_function(.5)
 
+  def test_stable(self):
+    # An extreme distribution that has probability 1 at 0 and probability 0
+    # otherwise.
+    dist = self.dist_cls(loc=[0, 0], scale=[0, 0], weight=[0.5, 0.5])
+    with self.subTest("AtTheMode"):
+      self.assertAllClose(dist.prob([0]), [1.0])
+    with self.subTest("NotAtTheMode"):
+      self.assertAllClose(dist.prob([1]), [0.0])
+
 
 class NoisyNormalMixtureTest(MixtureTest, tf.test.TestCase):