From ac993063900eb7ba9a6a6f4f282e55811ef3fec8 Mon Sep 17 00:00:00 2001
From: Zolisa Bleki <44142765+zoj613@users.noreply.github.com>
Date: Thu, 27 Jan 2022 11:12:06 +0200
Subject: [PATCH] Add DiscreteWeibull moment

---
 pymc/distributions/discrete.py           |  6 ++++++
 pymc/tests/test_distributions_moments.py | 22 +++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py
index 564b7c6ed0..021f430fc0 100644
--- a/pymc/distributions/discrete.py
+++ b/pymc/distributions/discrete.py
@@ -490,6 +490,12 @@ def dist(cls, q, beta, *args, **kwargs):
         beta = at.as_tensor_variable(floatX(beta))
         return super().dist([q, beta], **kwargs)
 
+    def get_moment(rv, size, q, beta):
+        median = at.power(at.log(0.5) / at.log(q), 1 / beta) - 1
+        if not rv_size_is_none(size):
+            median = at.full(size, median)
+        return median
+
     def logp(value, q, beta):
         r"""
         Calculate log-probability of DiscreteWeibull distribution at specified value.
diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py
index e87459f68a..5969bc9c26 100644
--- a/pymc/tests/test_distributions_moments.py
+++ b/pymc/tests/test_distributions_moments.py
@@ -22,6 +22,7 @@
     DensityDist,
     Dirichlet,
     DiscreteUniform,
+    DiscreteWeibull,
     ExGaussian,
     Exponential,
     Flat,
@@ -110,7 +111,6 @@ def test_all_distributions_have_moments():
 
     # Distributions that have been refactored but don't yet have moments
     not_implemented |= {
-        dist_module.discrete.DiscreteWeibull,
         dist_module.multivariate.DirichletMultinomial,
         dist_module.multivariate.Wishart,
     }
@@ -751,6 +751,26 @@ def test_discrete_uniform_moment(lower, upper, size, expected):
         DiscreteUniform("x", lower=lower, upper=upper, size=size)
 
 
+@pytest.mark.parametrize(
+    "q, beta, size, expected",
+    [
+        (0.5, 0.5, None, 0),
+        (0.6, 0.1, 5, (20,) * 5),
+        (np.linspace(0.25, 0.99, 4), 0.42, None, [0, 0, 6, 23862]),
+        (
+            np.linspace(0.5, 0.99, 3),
+            [[1, 1.25, 1.75], [1.25, 0.75, 0.5]],
+            None,
+            [[0, 0, 10], [0, 2, 4755]],
+        ),
+    ],
+)
+def test_discrete_weibull_moment(q, beta, size, expected):
+    with Model() as model:
+        DiscreteWeibull("x", q=q, beta=beta, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.parametrize(
     "a, size, expected",
     [