Fix numeric error for high-success BetaBinomial.

Specifically, BetaBinomial(concentration0=small, total_counts=n).log_prob(n)
used to incur a catastrophic cancellation between the total counts and
the count of successes, which lost accuracy in the concentration0
parameter.  By computing the count of failures before adding the
concentration we cover the case when there are zero failures without
losing precision.

PiperOrigin-RevId: 384326009

Author: axch

tensorflow_probability/python/distributions/beta_binomial.py

@@ -269,7 +269,7 @@ def _sample_n(self, n, seed=None):
   def _log_prob(self, counts):
     n, c1, c0 = self._params_list_as_tensors()
     return (_log_combinations(n, counts)
-            + tfp_math.lbeta(c1 + counts, n + c0 - counts)
+            + tfp_math.lbeta(c1 + counts, (n - counts) + c0)
             - tfp_math.lbeta(c1, c0))
 
   @distribution_util.AppendDocstring(_beta_binomial_sample_note)

tensorflow_probability/python/distributions/beta_binomial_test.py

@@ -149,6 +149,12 @@ def testSampleCornerConcentrations(self):
                          total_count=50.)
     self.assertAllEqual(d.sample(10, seed=seed_stream()), [[0, 50]] * 10)
 
+  def testLogProbCornerCase(self):
+    d = tfd.BetaBinomial(
+        concentration0=2e-8, concentration1=1.0, total_count=1.0)
+    # Numerically 0 is better than numerically +inf
+    self.assertAllEqual(d.log_prob(1.0), 0.0)
+
   def testEmpiricalCdfAgainstDirichletMultinomial(self):
     # This test is too slow for Eager mode.
     if tf.executing_eagerly():