Add MvNormal moment

pymc/distributions/multivariate.py

@@ -228,6 +228,13 @@ def dist(cls, mu, cov=None, tau=None, chol=None, lower=True, **kwargs):
         cov = quaddist_matrix(cov, chol, tau, lower)
         return super().dist([mu, cov], **kwargs)
 
+    def get_moment(rv, size, mu, cov):
+        moment = mu
+        if not rv_size_is_none(size):
+            moment_size = at.concatenate([size, mu.shape])
+            moment = at.full(moment_size, mu)
+        return moment
+
     def logp(value, mu, cov):
         """
         Calculate log-probability of Multivariate Normal distribution

pymc/tests/test_distributions_moments.py

@@ -43,6 +43,7 @@
     ZeroInflatedBinomial,
     ZeroInflatedPoisson,
 )
+from pymc.distributions.multivariate import MvNormal
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import Model
@@ -748,3 +749,21 @@ def test_categorical_moment(p, size, expected):
     with Model() as model:
         Categorical("x", p=p, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, cov, size, expected",
+    [
+        (np.ones(1), np.identity(1), None, np.ones(1)),
+        (np.ones(10), np.identity(10), None, np.ones(10)),
+        (np.ones(2), np.identity(2), 4, np.ones((4, 2))),
+        (np.ones(2), np.identity(2), (4, 2), np.ones((4, 2, 2))),
+        (np.ones((2, 2)), np.identity(2), None, np.ones((2, 2))),
+        (np.ones((2, 2)), np.identity(2), 4, np.ones((4, 2, 2))),
+        (np.ones((2, 2)), np.identity(2), (4, 2), np.ones((4, 2, 2, 2))),
+    ],
+)
+def test_mv_normal_moment(mu, cov, size, expected):
+    with Model() as model:
+        MvNormal("x", mu=mu, cov=cov, size=size)
+    assert_moment_is_expected(model, expected)