Add Gamma moment

Author: lucianopaz

pymc/distributions/continuous.py

@@ -2226,6 +2226,13 @@ def get_alpha_beta(cls, alpha=None, beta=None, mu=None, sigma=None):
 
         return alpha, beta
 
+    def get_moment(rv, size, alpha, inv_beta):
+        # The Aesara `GammaRV` `Op` inverts the `beta` parameter itself
+        mean = alpha * inv_beta
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logcdf(value, alpha, inv_beta):
         """
         Compute the log of the cumulative distribution function for Gamma distribution

pymc/tests/test_distributions_moments.py

@@ -6,6 +6,7 @@
     Beta,
     Cauchy,
     Exponential,
+    Gamma,
     HalfCauchy,
     HalfNormal,
     Kumaraswamy,
@@ -282,3 +283,23 @@ def test_halfcauchy_moment(beta, size, expected):
     with Model() as model:
         HalfCauchy("x", beta=beta, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "alpha, beta, size, expected",
+    [
+        (1, 1, None, 1),
+        (1, 1, 5, np.full(5, 1)),
+        (np.arange(1, 6), 1, None, np.arange(1, 6)),
+        (
+            np.arange(1, 6),
+            2 * np.arange(1, 6),
+            (2, 5),
+            np.full((2, 5), 0.5),
+        ),
+    ],
+)
+def test_gamma_moment(alpha, beta, size, expected):
+    with Model() as model:
+        Gamma("x", alpha=alpha, beta=beta, size=size)
+    assert_moment_is_expected(model, expected)