Added moment for geometric distribution along with test

Author: sagartomar

pymc/distributions/discrete.py

@@ -819,6 +819,12 @@ def dist(cls, p, *args, **kwargs):
         p = at.as_tensor_variable(floatX(p))
         return super().dist([p], *args, **kwargs)
 
+    def get_moment(rv, size, p):
+        mean = at.round(1.0 / p)
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logp(value, p):
         r"""
         Calculate log-probability of Geometric distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -12,6 +12,7 @@
     Constant,
     Exponential,
     Gamma,
+    Geometric,
     HalfCauchy,
     HalfNormal,
     HalfStudentT,
@@ -482,3 +483,18 @@ def test_logistic_moment(mu, s, size, expected):
     with Model() as model:
         Logistic("x", mu=mu, s=s, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "p, size, expected",
+    [
+        (0.5, None, 2),
+        (0.2, 5, 5 * np.ones(5)),
+        (np.linspace(0.25, 1, 4), None, [4, 2, 1, 1]),
+        (np.linspace(0.25, 1, 4), (2, 4), np.full((2, 4), [4, 2, 1, 1])),
+    ],
+)
+def test_geometric_moment(p, size, expected):
+    with Model() as model:
+        Geometric("x", p=p, size=size)
+    assert_moment_is_expected(model, expected)