Fix MatrixNormal.random

Author: Sayam753

pymc3/distributions/multivariate.py

@@ -1445,7 +1445,7 @@ class MatrixNormal(Continuous):
 
     .. math::
        f(x \mid \mu, U, V) =
-           \frac{1}{(2\pi |U|^n |V|^m)^{1/2}}
+           \frac{1}{(2\pi^{m n} |U|^n |V|^m)^{1/2}}
            \exp\left\{
                 -\frac{1}{2} \mathrm{Tr}[ V^{-1} (x-\mu)^{\prime} U^{-1} (x-\mu)]
             \right\}
@@ -1637,27 +1637,21 @@ def random(self, point=None, size=None):
         mu, colchol, rowchol = draw_values(
             [self.mu, self.colchol_cov, self.rowchol_cov], point=point, size=size
         )
-        if size is None:
-            size = ()
-        if size in (None, ()):
-            standard_normal = np.random.standard_normal((self.shape[0], colchol.shape[-1]))
-            samples = mu + np.matmul(rowchol, np.matmul(standard_normal, colchol.T))
-        else:
-            samples = []
-            size = tuple(np.atleast_1d(size))
-            if mu.shape == tuple(self.shape):
-                for _ in range(np.prod(size)):
-                    standard_normal = np.random.standard_normal((self.shape[0], colchol.shape[-1]))
-                    samples.append(mu + np.matmul(rowchol, np.matmul(standard_normal, colchol.T)))
-            else:
-                for j in range(np.prod(size)):
-                    standard_normal = np.random.standard_normal(
-                        (self.shape[0], colchol[j].shape[-1])
-                    )
-                    samples.append(
-                        mu[j] + np.matmul(rowchol[j], np.matmul(standard_normal, colchol[j].T))
-                    )
-            samples = np.array(samples).reshape(size + tuple(self.shape))
+        size = to_tuple(size)
+        dist_shape = to_tuple(self.shape)
+        output_shape = size + dist_shape
+
+        # Broadcasting all parameters
+        mu = broadcast_dist_samples_to(to_shape=output_shape, samples=[mu], size=size)[0]
+        rowchol = np.broadcast_to(rowchol, shape=size + rowchol.shape[-2:])
+
+        colchol = np.broadcast_to(colchol, shape=size + colchol.shape[-2:])
+        perm = np.arange(len(output_shape))
+        perm[-2:] = [perm[-1], perm[-2]]
+        colchol = np.transpose(colchol, perm)
+
+        standard_normal = np.random.standard_normal(output_shape)
+        samples = mu + np.matmul(rowchol, np.matmul(standard_normal, colchol))
         return samples
 
     def _trquaddist(self, value):

pymc3/tests/test_distributions_random.py

@@ -1675,3 +1675,21 @@ def test_issue_3706(self):
             prior_pred = pm.sample_prior_predictive(1)
 
         assert prior_pred["X"].shape == (1, N, 2)
+
+
+def test_issue_3585():
+    K = 3
+    D = 15
+    mu_0 = np.zeros((D, K))
+    lambd = 1.0
+    with pm.Model() as model:
+        sd_dist = pm.HalfCauchy.dist(beta=2.5)
+        packedL = pm.LKJCholeskyCov(f"packedL", eta=2, n=D, sd_dist=sd_dist)
+        L = pm.expand_packed_triangular(D, packedL, lower=True)
+        Sigma = pm.Deterministic(f"Sigma", L.dot(L.T))  # D x D covariance
+        mu = pm.MatrixNormal(
+            f"mu", mu=mu_0, rowcov=(1 / lambd) * Sigma, colcov=np.eye(K), shape=(D, K)
+        )
+        prior = pm.sample_prior_predictive(2)
+
+    assert prior["mu"].shape == (2, D, K)