Add MvNormal moment

pymc/distributions/multivariate.py

@@ -21,6 +21,7 @@
 
 import aesara
 import aesara.tensor as at
+from aesara.tensor.extra_ops import broadcast_shape
 import numpy as np
 import scipy
 
@@ -46,7 +47,11 @@
 from pymc.distributions.continuous import ChiSquared, Normal, assert_negative_support
 from pymc.distributions.dist_math import bound, factln, logpow, multigammaln
 from pymc.distributions.distribution import Continuous, Discrete
-from pymc.distributions.shape_utils import broadcast_dist_samples_to, to_tuple
+from pymc.distributions.shape_utils import (
+    broadcast_dist_samples_to, 
+    rv_size_is_none, 
+    to_tuple
+)
 from pymc.math import kron_diag, kron_dot
 
 __all__ = [
@@ -224,6 +229,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):
+            m_size = at.concatenate([size, mu.shape])
+            moment = at.full(m_size, mu)
+        return moment
+
     def logp(value, mu, cov):
         """
         Calculate log-probability of Multivariate Normal distribution

pymc/tests/test_distributions_moments.py

@@ -36,6 +36,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
@@ -595,3 +596,18 @@ def test_discrete_uniform_moment(lower, upper, size, expected):
     with Model() as model:
         DiscreteUniform("x", lower=lower, upper=upper, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, cov, size, expected",
+    [
+        (np.array([1.]), np.array([[1.]]), None, np.array([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))),
+    ]
+)
+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)