diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index d3bfe644c4..8578ab5a95 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -2628,16 +2628,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. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 36fe591e69..f41f9f63eb 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -11,6 +11,7 @@ Gamma, HalfCauchy, HalfNormal, + HalfStudentT, Kumaraswamy, Laplace, LogNormal, @@ -127,6 +128,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",