adds explicit rounding to get_moment

Author: morganstrom

pymc/distributions/discrete.py

@@ -240,7 +240,7 @@ def dist(cls, alpha, beta, n, *args, **kwargs):
         return super().dist([n, alpha, beta], **kwargs)
 
     def get_moment(rv, size, n, alpha, beta):
-        mean = (n * alpha) / (alpha + beta)
+        mean = at.round((n * alpha) / (alpha + beta))
         if not rv_size_is_none(size):
             mean = at.full(size, mean)
         return mean

pymc/tests/test_distributions_moments.py

@@ -215,8 +215,8 @@ def test_beta_moment(alpha, beta, size, expected):
     [
         (10, 1, 1, None, 5),
         (10, 1, 1, 5, np.full(5, 5)),
-        (10, 1, np.arange(1, 6), None, np.floor(10 / np.arange(2, 7))),
-        (10, 1, np.arange(1, 6), (2, 5), np.full((2, 5), np.floor(10 / np.arange(2, 7)))),
+        (10, 1, np.arange(1, 6), None, np.round(10 / np.arange(2, 7))),
+        (10, 1, np.arange(1, 6), (2, 5), np.full((2, 5), np.round(10 / np.arange(2, 7)))),
     ],
 )
 def test_beta_binomial_moment(alpha, beta, n, size, expected):