Make discrete moments robust against int-overflows

Author: michaelosthege

pymc/distributions/discrete.py

@@ -878,7 +878,11 @@ def dist(cls, N, k, n, *args, **kwargs):
         return super().dist([good, bad, n], *args, **kwargs)
 
     def moment(rv, size, good, bad, n):
-        N, k = good + bad, good
+        # Cast to float because the intX can be int8
+        # which could trigger an integer overflow below.
+        n = floatX(n)
+        k = floatX(good)
+        N = k + floatX(bad)
         mode = at.floor((n + 1) * (k + 1) / (N + 2))
         if not rv_size_is_none(size):
             mode = at.full(size, mode)
@@ -1014,6 +1018,8 @@ def dist(cls, lower, upper, *args, **kwargs):
         return super().dist([lower, upper], **kwargs)
 
     def moment(rv, size, lower, upper):
+        upper = floatX(upper)
+        lower = floatX(lower)
         mode = at.maximum(at.floor((upper + lower) / 2.0), lower)
         if not rv_size_is_none(size):
             mode = at.full(size, mode)

pymc/tests/distributions/util.py

@@ -579,7 +579,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True):
 
     assert moment.shape == expected.shape
     assert expected.shape == random_draw.shape
-    assert np.allclose(moment, expected)
+    np.testing.assert_allclose(moment, expected, atol=1e-10)
 
     if check_finite_logp:
         logp_moment = (