Merge branch 'main' into main

Author: shivam15s

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/distributions/discrete.py

@@ -1217,6 +1217,11 @@ def dist(cls, c, *args, **kwargs):
         c = at.as_tensor_variable(floatX(c))
         return super().dist([c], **kwargs)
 
+    def get_moment(rv, size, c):
+        if not rv_size_is_none(size):
+            c = at.full(size, c)
+        return c
+
     def logp(value, c):
         r"""
         Calculate log-probability of Constant distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -9,10 +9,12 @@
     Binomial,
     Cauchy,
     ChiSquared,
+    Constant,
     Exponential,
     Gamma,
     HalfCauchy,
     HalfNormal,
+    HalfStudentT,
     Kumaraswamy,
     Laplace,
     Logistic,
@@ -131,6 +133,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",
@@ -385,6 +402,20 @@ def test_poisson_moment(mu, size, expected):
     assert_moment_is_expected(model, expected)
 
 
+@pytest.mark.parametrize(
+    "c, size, expected",
+    [
+        (1, None, 1),
+        (1, 5, np.full(5, 1)),
+        (np.arange(1, 6), None, np.arange(1, 6)),
+    ],
+)
+def test_constant_moment(c, size, expected):
+    with Model() as model:
+        Constant("x", c=c, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.parametrize(
     "mu, s, size, expected",
     [
@@ -403,3 +434,4 @@ 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)
+