Add HalfStudentT moment (#5152)

Co-authored-by: Ricardo Vieira <[email protected]>

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

Author: 3 people

pymc/distributions/continuous.py

@@ -2634,16 +2634,18 @@ def dist(cls, nu=1, sigma=None, lam=None, sd=None, *args, **kwargs):
         lam, sigma = get_tau_sigma(lam, sigma)
         sigma = at.as_tensor_variable(sigma)
 
-        # mode = at.as_tensor_variable(0)
-        # median = at.as_tensor_variable(sigma)
-        # sd = at.as_tensor_variable(sigma)
-
         assert_negative_support(nu, "nu", "HalfStudentT")
         assert_negative_support(lam, "lam", "HalfStudentT")
         assert_negative_support(sigma, "sigma", "HalfStudentT")
 
         return super().dist([nu, sigma], *args, **kwargs)
 
+    def get_moment(rv, size, nu, sigma):
+        sigma, _ = at.broadcast_arrays(sigma, nu)
+        if not rv_size_is_none(size):
+            sigma = at.full(size, sigma)
+        return sigma
+
     def logp(value, nu, sigma):
         """
         Calculate log-probability of HalfStudentT distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -13,6 +13,7 @@
     Gamma,
     HalfCauchy,
     HalfNormal,
+    HalfStudentT,
     Kumaraswamy,
     Laplace,
     LogNormal,
@@ -130,6 +131,21 @@ def test_halfnormal_moment(sigma, size, expected):
     assert_moment_is_expected(model, expected)
 
 
+@pytest.mark.parametrize(
+    "nu, sigma, size, expected",
+    [
+        (1, 1, None, 1),
+        (1, 1, 5, np.ones(5)),
+        (1, np.arange(5), (2, 5), np.full((2, 5), np.arange(5))),
+        (np.arange(1, 6), 1, None, np.full(5, 1)),
+    ],
+)
+def test_halfstudentt_moment(nu, sigma, size, expected):
+    with Model() as model:
+        HalfStudentT("x", nu=nu, sigma=sigma, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.skip(reason="aeppl interval transform fails when both edges are None")
 @pytest.mark.parametrize(
     "mu, sigma, lower, upper, size, expected",