Allow for scalar or size 1 mu in MvNormal and MvStudentT

Author: ricardoV94

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

@@ -2945,7 +2945,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))",
@@ -2965,7 +2965,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

@@ -950,7 +950,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)
 
 

pymc/tests/test_distributions_random.py

@@ -1050,8 +1050,32 @@ class TestMvNormalCov(BaseTestDistribution):
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
         "check_rv_size",
+        "check_mu_broadcast_helper",
     ]
 
+    def check_mu_broadcast_helper(self):
+        """Test that mu is broadcasted to the shape of cov"""
+        x = pm.MvNormal.dist(mu=1, cov=np.eye(3))
+        mu = x.owner.inputs[3]
+        assert mu.eval().shape == (3,)
+
+        x = pm.MvNormal.dist(mu=np.ones(1), cov=np.eye(3))
+        mu = x.owner.inputs[3]
+        assert mu.eval().shape == (3,)
+
+        x = pm.MvNormal.dist(mu=np.ones((1, 1)), cov=np.eye(3))
+        mu = x.owner.inputs[3]
+        assert mu.eval().shape == (1, 3)
+
+        x = pm.MvNormal.dist(mu=np.ones((10, 1)), cov=np.eye(3))
+        mu = x.owner.inputs[3]
+        assert mu.eval().shape == (10, 3)
+
+        # Cov is artificually limited to being 2D
+        # x = pm.MvNormal.dist(mu=np.ones((10, 1)), cov=np.full((2, 3, 3), np.eye(3)))
+        # mu = x.owner.inputs[3]
+        # assert mu.eval().shape == (10, 2, 3)
+
 
 class TestMvNormalChol(BaseTestDistribution):
     pymc_dist = pm.MvNormal
@@ -1111,6 +1135,7 @@ def mvstudentt_rng_fn(self, size, nu, mu, cov, rng):
         "check_pymc_draws_match_reference",
         "check_rv_size",
         "check_errors",
+        "check_mu_broadcast_helper",
     ]
 
     def check_errors(self):
@@ -1124,6 +1149,29 @@ def check_errors(self):
                     cov=np.full((2, 2), np.ones(2)),
                 )
 
+    def check_mu_broadcast_helper(self):
+        """Test that mu is broadcasted to the shape of cov"""
+        x = pm.MvStudentT.dist(nu=4, mu=1, cov=np.eye(3))
+        mu = x.owner.inputs[4]
+        assert mu.eval().shape == (3,)
+
+        x = pm.MvStudentT.dist(nu=4, mu=np.ones(1), cov=np.eye(3))
+        mu = x.owner.inputs[4]
+        assert mu.eval().shape == (3,)
+
+        x = pm.MvStudentT.dist(nu=4, mu=np.ones((1, 1)), cov=np.eye(3))
+        mu = x.owner.inputs[4]
+        assert mu.eval().shape == (1, 3)
+
+        x = pm.MvStudentT.dist(nu=4, mu=np.ones((10, 1)), cov=np.eye(3))
+        mu = x.owner.inputs[4]
+        assert mu.eval().shape == (10, 3)
+
+        # Cov is artificually limited to being 2D
+        # x = pm.MvStudentT.dist(nu=4, mu=np.ones((10, 1)), cov=np.full((2, 3, 3), np.eye(3)))
+        # mu = x.owner.inputs[4]
+        # assert mu.eval().shape == (10, 2, 3)
+
 
 class TestMvStudentTChol(BaseTestDistribution):
     pymc_dist = pm.MvStudentT