Stabilize Negative Binomial with large mu/alpha

pymc/dims/distributions/scalar.py

@@ -316,7 +316,7 @@ class NegativeBinomial(DimDistribution):
 
     @classmethod
     def dist(cls, mu=None, alpha=None, *, p=None, n=None, **kwargs):
-        n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n)
+        n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n, math=ptx.math)
         return super().dist([n, p], **kwargs)
 
 

pymc/distributions/discrete.py

@@ -699,7 +699,7 @@ def dist(cls, mu=None, alpha=None, p=None, n=None, *args, **kwargs):
         return super().dist([n, p], *args, **kwargs)
 
     @classmethod
-    def get_n_p(cls, mu=None, alpha=None, p=None, n=None):
+    def get_n_p(cls, mu=None, alpha=None, p=None, n=None, math=pt):
         if n is None:
             if alpha is not None:
                 n = alpha
@@ -710,7 +710,9 @@ def get_n_p(cls, mu=None, alpha=None, p=None, n=None):
 
         if p is None:
             if mu is not None:
-                p = n / (mu + n)
+                # n / (mu + n) in logit space, so a mu = exp(x) that overflows still
+                # yields finite logp and dlogp: log(mu) cancels back to x
+                p = math.sigmoid(math.log(n) - math.log(mu))
             else:
                 raise ValueError("Incompatible parametrization. Must specify either mu or p.")
         elif mu is not None:
@@ -725,7 +727,10 @@ def support_point(rv, size, n, p):
         return mu
 
     def logp(value, n, p):
-        mu = n * (1 - p) / p
+        # (1 - p) / p in log space: a p = sigmoid(w) that saturates at 1.0 carries no
+        # information, while the rewritten log terms still see w. Spelled log(1 - p)
+        # because the sigmoid stabilization rewrites do not recognize log1p(-p)
+        mu = n * pt.exp(pt.log(1 - p) - pt.log(p))
 
         # binomln subtracts gammaln(value + n) - gammaln(n), whose difference falls below
         # their shared ulp once n is large, so fall back on the Poisson(mu) limit there.

tests/distributions/test_discrete.py

@@ -33,7 +33,7 @@
 from pymc.exceptions import ImputationWarning
 from pymc.logprob.basic import icdf, logcdf, logp
 from pymc.logprob.utils import ParameterValueError
-from pymc.pytensorf import floatX
+from pymc.pytensorf import floatX, rewrite_pregrad
 from pymc.testing import (
     BaseTestDistributionRandom,
     Bool,
@@ -772,16 +772,43 @@ def test_negative_binomial_logp_large_n():
     np.testing.assert_allclose(logp_expr.eval({mu: 5.0, n: 1e18}), -1.9634457319257537)
     np.testing.assert_allclose(logp_expr.eval({mu: 5.0, n: 1e20}), -1.9634457319257537)
 
+    # a tiny mu saturates p = n / (mu + n) at 1.0, so the Poisson branch must recover
+    # log(mu) from the log-space terms for logp and dlogp to stay usable
+    a = pt.dscalar("a")
+    logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=1e12), 3)
+    np.testing.assert_allclose(logp_expr.eval({a: -300.0}), -901.7917594692281)
+
+    dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a)
+    np.testing.assert_allclose(dlogp_expr.eval({a: -300.0}), 3.0)
+
 
-@pytest.mark.xfail(
-    reason="Needs a log(a + exp(x)) -> log(a) + log1pexp(x - log(a)) stabilization in "
-    "PyTensor, since logp only ever sees p = n / (mu + n)"
-)
 def test_negative_binomial_logp_stable_when_mu_overflows():
+    """get_n_p builds p = sigmoid(log(n) - log(mu)), so with mu = exp(a) the log
+    cancels and the logp reduces to softplus terms in a, keeping logp and dlogp
+    finite when exp(a) overflows.
+
+    A constant alpha and a log-transformed alpha (the shape a positive alpha RV takes
+    in the logp graph) are both checked.
+    """
     a = pt.dscalar("a")
     logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=2.0), 3)
 
     np.testing.assert_allclose(logp_expr.eval({a: 710.0}), -1417.2274112777604)
+    np.testing.assert_allclose(logp_expr.eval({a: 5000.0}), -9997.22741127776)
+
+    dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a)
+    np.testing.assert_allclose(dlogp_expr.eval({a: 710.0}), -2.0)
+    np.testing.assert_allclose(dlogp_expr.eval({a: 5000.0}), -2.0)
+
+    b = pt.dscalar("b")
+    logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=pt.exp(b)), 3)
+
+    np.testing.assert_allclose(logp_expr.eval({a: 710.0, b: 0.7}), -1426.9536421881814)
+    np.testing.assert_allclose(logp_expr.eval({a: 5000.0, b: 0.7}), -10065.952757236526)
+
+    dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a)
+    np.testing.assert_allclose(dlogp_expr.eval({a: 710.0, b: 0.7}), -2.0137527074704766)
+    np.testing.assert_allclose(dlogp_expr.eval({a: 5000.0, b: 0.7}), -2.0137527074704766)
 
 
 class TestNegativeBinomialMuSigma(BaseTestDistributionRandom):

tests/distributions/test_mixture.py

@@ -1420,14 +1420,14 @@ def logcdf_fn(value, psi, mu):
 
     def test_zeroinflatednegativebinomial_logp(self):
         def logp_fn(value, psi, mu, alpha):
-            n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha)
+            n, p = alpha, alpha / (mu + alpha)
             if value == 0:
                 return np.log((1 - psi) * st.nbinom.pmf(0, n, p))
             else:
                 return np.log(psi * st.nbinom.pmf(value, n, p))
 
         def logcdf_fn(value, psi, mu, alpha):
-            n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha)
+            n, p = alpha, alpha / (mu + alpha)
             return np.log((1 - psi) + psi * st.nbinom.cdf(value, n, p))
 
         check_logp(
@@ -1537,10 +1537,12 @@ def test_zero_inflated_binomial_support_point(self, psi, n, p, size, expected):
             (0.2, 10, 4, 5, np.full(5, 2)),
             (
                 0.4,
-                np.arange(1, 5),
+                # keep the means off whole numbers, where the floor in the support
+                # point is sensitive to how p was built
+                np.arange(1, 5) + 0.5,
                 np.arange(2, 6),
                 None,
-                np.array([0, 1, 1, 2] if pytensor.config.floatX == "float64" else [0, 0, 1, 1]),
+                np.array([0, 1, 1, 2]),
             ),
             (
                 np.linspace(0.2, 0.6, 3),
@@ -1691,7 +1693,7 @@ def logp_fn(value, psi, mu):
 
     def test_hurdle_negativebinomial_logp(self):
         def logp_fn(value, psi, mu, alpha):
-            n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha)
+            n, p = alpha, alpha / (mu + alpha)
             if value == 0:
                 return np.log(1 - psi)
             else: