Add LKJCholeskyCov moment

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -1202,6 +1202,12 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs):
 
         return super().dist([n, eta, sd_dist], size=size, **kwargs)
 
+    def get_moment(rv, size, n, eta, sd_dists):
+        diag_idxs = (at.cumsum(at.arange(1, n + 1)) - 1).astype("int32")
+        moment = at.zeros_like(rv)
+        moment = at.set_subtensor(moment[..., diag_idxs], 1)
+        return moment
+
     def logp(value, n, eta, sd_dist):
         """
         Calculate log-probability of Covariance matrix with LKJ

pymc/tests/test_distributions_moments.py

@@ -39,6 +39,7 @@
     KroneckerNormal,
     Kumaraswamy,
     Laplace,
+    LKJCholeskyCov,
     LKJCorr,
     Logistic,
     LogitNormal,
@@ -1439,3 +1440,24 @@ def test_lkjcorr_moment(n, eta, size, expected):
     with Model() as model:
         LKJCorr("x", n=n, eta=eta, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "n, eta, size, expected",
+    [
+        (3, 1, None, np.array([1, 0, 1, 0, 0, 1])),
+        (4, 1, None, np.array([1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0])),
+        (3, 1, 1, np.array([[1, 0, 1, 0, 0, 1]])),
+        (
+            4,
+            1,
+            (2, 3),
+            np.full((2, 3, 10), np.array([1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0])),
+        ),
+    ],
+)
+def test_lkjcholeskycov_moment(n, eta, size, expected):
+    with Model() as model:
+        sd_dist = pm.Exponential.dist(1, size=(*to_tuple(size), n))
+        LKJCholeskyCov("x", n=n, eta=eta, sd_dist=sd_dist, size=size, compute_corr=False)
+    assert_moment_is_expected(model, expected, check_finite_logp=size is None)