From 684f338b85db37e0a15114c35bc0b4674dcf0b5b Mon Sep 17 00:00:00 2001
From: kc611 <ckaustubhm06@gmail.com>
Date: Wed, 23 Jun 2021 17:24:32 +0530
Subject: [PATCH] Tested each dimension seperately in MatrixNormal tests

---
 pymc3/tests/test_distributions_random.py | 40 ++++++++++++++++++------
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py
index 37d6aaefaed..9171c731196 100644
--- a/pymc3/tests/test_distributions_random.py
+++ b/pymc3/tests/test_distributions_random.py
@@ -41,10 +41,8 @@
 from pymc3.tests.test_distributions import (
     Domain,
     Nat,
-    PdMatrix,
     R,
     RandomPdMatrix,
-    RealMatrix,
     Rplus,
     Rplusbig,
     Simplex,
@@ -1352,16 +1350,40 @@ class TestMatrixNormal(BaseTestDistribution):
     tests_to_run = ["check_rv_size", "check_pymc_params_match_rv_op", "test_matrix_normal"]
 
     def test_matrix_normal(self):
+        delta = 0.05  # limit for KS p-value
+        n_fails = 100  # Allows the KS fails a certain number of times
+        size = (100,)
+
         def ref_rand(size, mu, rowcov, colcov):
             return st.matrix_normal.rvs(mean=mu, rowcov=rowcov, colcov=colcov, size=size)
 
-        pymc3_random(
-            pm.MatrixNormal,
-            {"mu": RealMatrix(3, 3), "rowcov": PdMatrix(3), "colcov": PdMatrix(3)},
-            size=100,
-            valuedomain=RealMatrix(3, 3),
-            ref_rand=ref_rand,
-        )
+        with pm.Model(rng_seeder=1):
+            matrixnormal = pm.MatrixNormal(
+                "mvnormal",
+                mu=np.random.random((3, 3)),
+                rowcov=np.eye(3),
+                colcov=np.eye(3),
+                size=size,
+            )
+            check = pm.sample_prior_predictive(n_fails)
+
+        ref_smp = ref_rand(size[0], mu=np.random.random((3, 3)), rowcov=np.eye(3), colcov=np.eye(3))
+
+        p, f = delta, n_fails
+        while p <= delta and f > 0:
+            matrixnormal_smp = check["mvnormal"][f - 1, :, :]
+            curr_ref_smp = ref_smp[f - 1, :, :]
+
+            p = min(
+                st.ks_2samp(
+                    np.atleast_1d(matrixnormal_smp[..., idx]).flatten(),
+                    np.atleast_1d(curr_ref_smp[..., idx]).flatten(),
+                )[1]
+                for idx in range(matrixnormal_smp.shape[-1])
+            )
+            f -= 1
+
+        assert p > delta
 
 
 class TestInterpolated(BaseTestDistribution):