Add wald and pareto moments (#5161)

Co-authored-by: Alihan Zihna <[email protected]>

Author: azihna

pymc/distributions/continuous.py

@@ -995,6 +995,12 @@ def dist(
 
         return super().dist([mu, lam, alpha], **kwargs)
 
+    def get_moment(rv, size, mu, lam, alpha):
+        mu, _, _ = at.broadcast_arrays(mu, lam, alpha)
+        if not rv_size_is_none(size):
+            mu = at.full(size, mu)
+        return mu
+
     @staticmethod
     def get_mu_lam_phi(
         mu: Optional[float], lam: Optional[float], phi: Optional[float]
@@ -1943,8 +1949,11 @@ def dist(
 
         return super().dist([alpha, m], **kwargs)
 
-    def _distr_parameters_for_repr(self):
-        return ["alpha", "m"]
+    def get_moment(rv, size, alpha, m):
+        median = m * 2 ** (1 / alpha)
+        if not rv_size_is_none(size):
+            median = at.full(size, median)
+        return median
 
     def logcdf(
         value: Union[float, np.ndarray, TensorVariable],

pymc/tests/test_distributions_moments.py

@@ -26,10 +26,12 @@
     LogNormal,
     NegativeBinomial,
     Normal,
+    Pareto,
     Poisson,
     StudentT,
     TruncatedNormal,
     Uniform,
+    Wald,
     Weibull,
     ZeroInflatedBinomial,
     ZeroInflatedPoisson,
@@ -360,6 +362,42 @@ def test_gamma_moment(alpha, beta, size, expected):
     assert_moment_is_expected(model, expected)
 
 
+@pytest.mark.parametrize(
+    "alpha, m, size, expected",
+    [
+        (2, 1, None, 1 * 2 ** (1 / 2)),
+        (2, 1, 5, np.full(5, 1 * 2 ** (1 / 2))),
+        (np.arange(2, 7), np.arange(1, 6), None, np.arange(1, 6) * 2 ** (1 / np.arange(2, 7))),
+        (
+            np.arange(2, 7),
+            np.arange(1, 6),
+            (2, 5),
+            np.full((2, 5), np.arange(1, 6) * 2 ** (1 / np.arange(2, 7))),
+        ),
+    ],
+)
+def test_pareto_moment(alpha, m, size, expected):
+    with Model() as model:
+        Pareto("x", alpha=alpha, m=m, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, lam, phi, size, expected",
+    [
+        (2, None, None, None, 2),
+        (None, 1, 1, 5, np.full(5, 1)),
+        (1, None, np.ones(5), None, np.full(5, 1)),
+        (3, np.full(5, 2), None, None, np.full(5, 3)),
+        (np.arange(1, 6), None, np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(1, 6))),
+    ],
+)
+def test_wald_moment(mu, lam, phi, size, expected):
+    with Model() as model:
+        Wald("x", mu=mu, lam=lam, phi=phi, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.parametrize(
     "alpha, beta, size, expected",
     [