Allow for scalar or size 1 mu in MvNormal and MvStudentT

pymc/distributions/multivariate.py

@@ -233,6 +233,8 @@ class MvNormal(Continuous):
     def dist(cls, mu, cov=None, tau=None, chol=None, lower=True, **kwargs):
         mu = at.as_tensor_variable(mu)
         cov = quaddist_matrix(cov, chol, tau, lower)
+        # Aesara is stricter about the shape of mu, than PyMC used to be
+        mu = at.broadcast_arrays(mu, cov[..., -1])[0]
         return super().dist([mu, cov], **kwargs)
 
     def get_moment(rv, size, mu, cov):
@@ -362,6 +364,8 @@ def dist(cls, nu, Sigma=None, mu=None, cov=None, tau=None, chol=None, lower=True
         nu = at.as_tensor_variable(floatX(nu))
         mu = at.as_tensor_variable(floatX(mu))
         cov = quaddist_matrix(cov, chol, tau, lower)
+        # Aesara is stricter about the shape of mu, than PyMC used to be
+        mu = at.broadcast_arrays(mu, cov[..., -1])[0]
         assert_negative_support(nu, "nu", "MvStudentT")
         return super().dist([nu, mu, cov], **kwargs)
 

pymc/tests/test_distributions.py

@@ -2903,7 +2903,7 @@ def setup_class(self):
                 r"sigma ~ N**+(0, 1)",
                 r"mu ~ Deterministic(f(beta, alpha))",
                 r"beta ~ N(0, 10)",
-                r"Z ~ N(<constant>, f())",
+                r"Z ~ N(f(), f())",
                 r"nb_with_p_n ~ NB(10, nbp)",
                 r"Y_obs ~ N(mu, sigma)",
                 r"pot ~ Potential(f(beta, alpha))",
@@ -2923,7 +2923,7 @@ def setup_class(self):
                 r"$\text{sigma} \sim \operatorname{N^{+}}(0,~1)$",
                 r"$\text{mu} \sim \operatorname{Deterministic}(f(\text{beta},~\text{alpha}))$",
                 r"$\text{beta} \sim \operatorname{N}(0,~10)$",
-                r"$\text{Z} \sim \operatorname{N}(\text{<constant>},~f())$",
+                r"$\text{Z} \sim \operatorname{N}(f(),~f())$",
                 r"$\text{nb_with_p_n} \sim \operatorname{NB}(10,~\text{nbp})$",
                 r"$\text{Y_obs} \sim \operatorname{N}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",

pymc/tests/test_distributions_moments.py

@@ -929,7 +929,7 @@ def test_mv_normal_moment(mu, cov, size, expected):
     with Model() as model:
         x = MvNormal("x", mu=mu, cov=cov, size=size)
 
-    # MvNormal logp is only impemented for up to 2D variables
+    # MvNormal logp is only implemented for up to 2D variables
     assert_moment_is_expected(model, expected, check_finite_logp=x.ndim < 3)