From fe1bf1bc65f291f0518bf041807ed4642adb38f9 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Mon, 2 Nov 2020 22:20:59 +0530
Subject: [PATCH 01/13] Fixed MvNormal.random method

---
 pymc3/distributions/multivariate.py | 74 ++++++++++++-----------------
 1 file changed, 31 insertions(+), 43 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 79831784b0..effe6ecb72 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -31,7 +31,14 @@
 
 from pymc3.theanof import floatX
 from . import transforms
-from .distribution import Continuous, Discrete, draw_values, generate_samples, _DrawValuesContext
+from .distribution import (
+    Continuous,
+    Discrete,
+    draw_values,
+    generate_samples,
+    _DrawValuesContext,
+    is_fast_drawable,
+)
 from ..model import Deterministic
 from .continuous import ChiSquared, Normal
 from .special import gammaln, multigammaln
@@ -250,56 +257,37 @@ def random(self, point=None, size=None):
         -------
         array
         """
-        if size is None:
-            size = tuple()
+        size = to_tuple(size)
+
+        param_attribute = "chol_cov" if self._cov_type == "chol" else self._cov_type
+        mu, param = draw_values([self.mu, getattr(self, param_attribute)], point=point, size=size)
+        output_shape = size + tuple(self.shape)
+        extra_dims = len(output_shape) - mu.ndim
+
+        if is_fast_drawable(self.mu):
+            mu = mu.reshape((1,) * extra_dims + mu.shape)
         else:
-            if not isinstance(size, tuple):
-                try:
-                    size = tuple(size)
-                except TypeError:
-                    size = (size,)
+            mu = mu.reshape(size + (1,) * extra_dims + mu.shape[len(size) :])
 
-        if self._cov_type == "cov":
-            mu, cov = draw_values([self.mu, self.cov], point=point, size=size)
-            if mu.shape[-1] != cov.shape[-1]:
-                raise ValueError("Shapes for mu and cov don't match")
+        mu = np.broadcast_to(mu, output_shape)
 
-            try:
-                dist = stats.multivariate_normal(mean=mu, cov=cov, allow_singular=True)
-            except ValueError:
-                size += (mu.shape[-1],)
-                return np.nan * np.zeros(size)
-            return dist.rvs(size)
-        elif self._cov_type == "chol":
-            mu, chol = draw_values([self.mu, self.chol_cov], point=point, size=size)
-            if size and mu.ndim == len(size) and mu.shape == size:
-                mu = mu[..., np.newaxis]
-            if mu.shape[-1] != chol.shape[-1] and mu.shape[-1] != 1:
-                raise ValueError("Shapes for mu and chol don't match")
-            broadcast_shape = np.broadcast(np.empty(mu.shape[:-1]), np.empty(chol.shape[:-2])).shape
-
-            mu = np.broadcast_to(mu, broadcast_shape + (chol.shape[-1],))
-            chol = np.broadcast_to(chol, broadcast_shape + chol.shape[-2:])
-            # If mu and chol were fixed by the point, only the standard normal
-            # should change
-            if mu.shape[: len(size)] != size:
-                std_norm_shape = size + mu.shape
-            else:
-                std_norm_shape = mu.shape
-            standard_normal = np.random.standard_normal(std_norm_shape)
+        if mu.shape[-1] != param.shape[-1]:
+            raise ValueError(f"Shapes for mu and {self._cov_type} don't match")
+
+        if self._cov_type == "cov":
+            chol = linalg.cholesky(param, lower=True)
+            standard_normal = np.random.standard_normal(output_shape)
             return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
+        elif self._cov_type == "chol":
+            standard_normal = np.random.standard_normal(output_shape)
+            return mu + np.einsum("...ij,...j->...i", param, standard_normal)
         else:
-            mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
-            if mu.shape[-1] != tau[0].shape[-1]:
-                raise ValueError("Shapes for mu and tau don't match")
-
-            size += (mu.shape[-1],)
             try:
-                chol = linalg.cholesky(tau, lower=True)
+                chol = linalg.cholesky(param, lower=True)
             except linalg.LinAlgError:
-                return np.nan * np.zeros(size)
+                return np.nan * np.zeros(output_shape)
 
-            standard_normal = np.random.standard_normal(size)
+            standard_normal = np.random.standard_normal(output_shape)
             transformed = linalg.solve_triangular(chol, standard_normal.T, lower=True)
             return mu + transformed.T
 

From 31d9ffb174f1dda20a692f148bcc44646354e31f Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Mon, 2 Nov 2020 22:35:42 +0530
Subject: [PATCH 02/13] Added some comments

---
 pymc3/distributions/multivariate.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index effe6ecb72..5861181922 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -261,9 +261,14 @@ def random(self, point=None, size=None):
 
         param_attribute = "chol_cov" if self._cov_type == "chol" else self._cov_type
         mu, param = draw_values([self.mu, getattr(self, param_attribute)], point=point, size=size)
+        # self.shape can be event or batch+event. So if it is given,
+        # deterministic nature of random method can be obtained by appending it to sample_shape.
         output_shape = size + tuple(self.shape)
         extra_dims = len(output_shape) - mu.ndim
 
+        # It was not a good idea to check mu.shape[:len(size)] == size,
+        # because it can get mixed among batch and event dimensions. Here, we explicitly chop off
+        # the size (sample_shape) and only broadcast batch and event dimensions.
         if is_fast_drawable(self.mu):
             mu = mu.reshape((1,) * extra_dims + mu.shape)
         else:
@@ -275,6 +280,7 @@ def random(self, point=None, size=None):
             raise ValueError(f"Shapes for mu and {self._cov_type} don't match")
 
         if self._cov_type == "cov":
+            # Calculate mu + L * samples, L = chol(cov)
             chol = linalg.cholesky(param, lower=True)
             standard_normal = np.random.standard_normal(output_shape)
             return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
@@ -282,6 +288,7 @@ def random(self, point=None, size=None):
             standard_normal = np.random.standard_normal(output_shape)
             return mu + np.einsum("...ij,...j->...i", param, standard_normal)
         else:
+            # Calculate mu + L^{-T} * samples, L = chol(tau)
             try:
                 chol = linalg.cholesky(param, lower=True)
             except linalg.LinAlgError:

From 7ded21c2913629a4ed99aaa61f30cb58b19afbbc Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Tue, 3 Nov 2020 14:19:02 +0530
Subject: [PATCH 03/13] Handled corner case when self.shape is not provided

---
 pymc3/distributions/multivariate.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 5861181922..84f56476fb 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -261,9 +261,18 @@ def random(self, point=None, size=None):
 
         param_attribute = "chol_cov" if self._cov_type == "chol" else self._cov_type
         mu, param = draw_values([self.mu, getattr(self, param_attribute)], point=point, size=size)
+        if tuple(self.shape):
+            dist_shape = tuple(self.shape)
+        else:
+            if is_fast_drawable(self.mu):
+                batch_shape = mu.shape[:-1]
+            else:
+                batch_shape = mu.shape[len(size) : -1]
+            dist_shape = batch_shape + param.shape[-1:]
+
         # self.shape can be event or batch+event. So if it is given,
         # deterministic nature of random method can be obtained by appending it to sample_shape.
-        output_shape = size + tuple(self.shape)
+        output_shape = size + dist_shape
         extra_dims = len(output_shape) - mu.ndim
 
         # It was not a good idea to check mu.shape[:len(size)] == size,

From e8ad24262d7732046fe76525a9d4b7b8e5b0834a Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Tue, 3 Nov 2020 14:25:22 +0530
Subject: [PATCH 04/13] Fixed comments

---
 pymc3/distributions/multivariate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 84f56476fb..eb09dd47fd 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -270,7 +270,7 @@ def random(self, point=None, size=None):
                 batch_shape = mu.shape[len(size) : -1]
             dist_shape = batch_shape + param.shape[-1:]
 
-        # self.shape can be event or batch+event. So if it is given,
+        # First, distribution shape (batch+event) is computed and then
         # deterministic nature of random method can be obtained by appending it to sample_shape.
         output_shape = size + dist_shape
         extra_dims = len(output_shape) - mu.ndim

From fd75e6ff583d85ada23020221a1d5f23a31577f4 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Sat, 7 Nov 2020 13:15:49 +0530
Subject: [PATCH 05/13] Considered the corner case of 'point' as well.

---
 pymc3/distributions/multivariate.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index eb09dd47fd..5682a532ef 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -264,7 +264,7 @@ def random(self, point=None, size=None):
         if tuple(self.shape):
             dist_shape = tuple(self.shape)
         else:
-            if is_fast_drawable(self.mu):
+            if is_fast_drawable(self.mu) or point:
                 batch_shape = mu.shape[:-1]
             else:
                 batch_shape = mu.shape[len(size) : -1]
@@ -278,7 +278,7 @@ def random(self, point=None, size=None):
         # It was not a good idea to check mu.shape[:len(size)] == size,
         # because it can get mixed among batch and event dimensions. Here, we explicitly chop off
         # the size (sample_shape) and only broadcast batch and event dimensions.
-        if is_fast_drawable(self.mu):
+        if is_fast_drawable(self.mu) or point:
             mu = mu.reshape((1,) * extra_dims + mu.shape)
         else:
             mu = mu.reshape(size + (1,) * extra_dims + mu.shape[len(size) :])

From b328931e1a62278b6588817afab6fd97e72613f6 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Mon, 2 Nov 2020 22:20:59 +0530
Subject: [PATCH 06/13] Fixed MvNormal.random method

---
 pymc3/distributions/multivariate.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 5682a532ef..69377af430 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -289,7 +289,6 @@ def random(self, point=None, size=None):
             raise ValueError(f"Shapes for mu and {self._cov_type} don't match")
 
         if self._cov_type == "cov":
-            # Calculate mu + L * samples, L = chol(cov)
             chol = linalg.cholesky(param, lower=True)
             standard_normal = np.random.standard_normal(output_shape)
             return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
@@ -297,7 +296,6 @@ def random(self, point=None, size=None):
             standard_normal = np.random.standard_normal(output_shape)
             return mu + np.einsum("...ij,...j->...i", param, standard_normal)
         else:
-            # Calculate mu + L^{-T} * samples, L = chol(tau)
             try:
                 chol = linalg.cholesky(param, lower=True)
             except linalg.LinAlgError:

From 894e529fdb17bbe933558a50d2a645464bb3f9e7 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Fri, 20 Nov 2020 20:53:25 +0530
Subject: [PATCH 07/13] Handled sample and batch dimensions in tau
 parametrization using numpy

Added batch dimensions to all parametrization
---
 pymc3/distributions/multivariate.py | 36 +++++++++++++++--------------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 69377af430..1ace806ffd 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -259,12 +259,16 @@ def random(self, point=None, size=None):
         """
         size = to_tuple(size)
 
-        param_attribute = "chol_cov" if self._cov_type == "chol" else self._cov_type
-        mu, param = draw_values([self.mu, getattr(self, param_attribute)], point=point, size=size)
+        param_attribute = getattr(self, "chol_cov" if self._cov_type == "chol" else self._cov_type)
+        mu, param = draw_values([self.mu, param_attribute], point=point, size=size)
+        check_fast_drawable_or_point = lambda param, point: is_fast_drawable(param) or (
+            point and hasattr(param, "model") and param.name in point
+        )
         if tuple(self.shape):
             dist_shape = tuple(self.shape)
+            batch_shape = dist_shape[:-1]
         else:
-            if is_fast_drawable(self.mu) or point:
+            if check_fast_drawable_or_point(self.mu, point):
                 batch_shape = mu.shape[:-1]
             else:
                 batch_shape = mu.shape[len(size) : -1]
@@ -278,32 +282,30 @@ def random(self, point=None, size=None):
         # It was not a good idea to check mu.shape[:len(size)] == size,
         # because it can get mixed among batch and event dimensions. Here, we explicitly chop off
         # the size (sample_shape) and only broadcast batch and event dimensions.
-        if is_fast_drawable(self.mu) or point:
+        if check_fast_drawable_or_point(self.mu, point):
             mu = mu.reshape((1,) * extra_dims + mu.shape)
         else:
             mu = mu.reshape(size + (1,) * extra_dims + mu.shape[len(size) :])
 
-        mu = np.broadcast_to(mu, output_shape)
+        # Adding batch dimensions to parametrization
+        if not check_fast_drawable_or_point(param_attribute, point):
+            param = param.reshape(size + (1,) * len(batch_shape) + param.shape[-2:])
 
+        mu = np.broadcast_to(mu, output_shape)
+        param = np.broadcast_to(param, output_shape + param.shape[-1:])
         if mu.shape[-1] != param.shape[-1]:
             raise ValueError(f"Shapes for mu and {self._cov_type} don't match")
 
         if self._cov_type == "cov":
-            chol = linalg.cholesky(param, lower=True)
-            standard_normal = np.random.standard_normal(output_shape)
-            return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
+            chol = np.linalg.cholesky(param)
         elif self._cov_type == "chol":
-            standard_normal = np.random.standard_normal(output_shape)
-            return mu + np.einsum("...ij,...j->...i", param, standard_normal)
+            chol = param
         else:
-            try:
-                chol = linalg.cholesky(param, lower=True)
-            except linalg.LinAlgError:
-                return np.nan * np.zeros(output_shape)
+            inverse = np.linalg.inv(param)
+            chol = np.linalg.cholesky(inverse)
 
-            standard_normal = np.random.standard_normal(output_shape)
-            transformed = linalg.solve_triangular(chol, standard_normal.T, lower=True)
-            return mu + transformed.T
+        standard_normal = np.random.standard_normal(output_shape)
+        return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
 
     def logp(self, value):
         """

From b52f1bb102ccbc79d174bb752b8692092ca21158 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Sat, 21 Nov 2020 16:46:09 +0530
Subject: [PATCH 08/13] Modified logic while inserting batch dimensions to
 parametrization

---
 pymc3/distributions/multivariate.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 1ace806ffd..ab53ec3630 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -262,8 +262,9 @@ def random(self, point=None, size=None):
         param_attribute = getattr(self, "chol_cov" if self._cov_type == "chol" else self._cov_type)
         mu, param = draw_values([self.mu, param_attribute], point=point, size=size)
         check_fast_drawable_or_point = lambda param, point: is_fast_drawable(param) or (
-            point and hasattr(param, "model") and param.name in point
+            point and param.name in point
         )
+
         if tuple(self.shape):
             dist_shape = tuple(self.shape)
             batch_shape = dist_shape[:-1]
@@ -288,7 +289,7 @@ def random(self, point=None, size=None):
             mu = mu.reshape(size + (1,) * extra_dims + mu.shape[len(size) :])
 
         # Adding batch dimensions to parametrization
-        if not check_fast_drawable_or_point(param_attribute, point):
+        if size and param.shape[:-2] == size:
             param = param.reshape(size + (1,) * len(batch_shape) + param.shape[-2:])
 
         mu = np.broadcast_to(mu, output_shape)

From c25a98837a346539b7b1bf6de9d467d9ed4245d7 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Fri, 27 Nov 2020 20:10:19 +0530
Subject: [PATCH 09/13] Used shapes_utils.broadcast_dist_samples_to function
 for broadcasting

---
 pymc3/distributions/multivariate.py | 56 ++++++++---------------------
 1 file changed, 15 insertions(+), 41 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index ab53ec3630..a998499cad 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -43,7 +43,7 @@
 from .continuous import ChiSquared, Normal
 from .special import gammaln, multigammaln
 from .dist_math import bound, logpow, factln
-from .shape_utils import to_tuple
+from .shape_utils import to_tuple, broadcast_dist_samples_to
 from ..math import kron_dot, kron_diag, kron_solve_lower, kronecker
 
 
@@ -261,50 +261,23 @@ def random(self, point=None, size=None):
 
         param_attribute = getattr(self, "chol_cov" if self._cov_type == "chol" else self._cov_type)
         mu, param = draw_values([self.mu, param_attribute], point=point, size=size)
-        check_fast_drawable_or_point = lambda param, point: is_fast_drawable(param) or (
-            point and param.name in point
-        )
-
-        if tuple(self.shape):
-            dist_shape = tuple(self.shape)
-            batch_shape = dist_shape[:-1]
-        else:
-            if check_fast_drawable_or_point(self.mu, point):
-                batch_shape = mu.shape[:-1]
-            else:
-                batch_shape = mu.shape[len(size) : -1]
-            dist_shape = batch_shape + param.shape[-1:]
 
-        # First, distribution shape (batch+event) is computed and then
-        # deterministic nature of random method can be obtained by appending it to sample_shape.
-        output_shape = size + dist_shape
-        extra_dims = len(output_shape) - mu.ndim
-
-        # It was not a good idea to check mu.shape[:len(size)] == size,
-        # because it can get mixed among batch and event dimensions. Here, we explicitly chop off
-        # the size (sample_shape) and only broadcast batch and event dimensions.
-        if check_fast_drawable_or_point(self.mu, point):
-            mu = mu.reshape((1,) * extra_dims + mu.shape)
-        else:
-            mu = mu.reshape(size + (1,) * extra_dims + mu.shape[len(size) :])
-
-        # Adding batch dimensions to parametrization
-        if size and param.shape[:-2] == size:
-            param = param.reshape(size + (1,) * len(batch_shape) + param.shape[-2:])
-
-        mu = np.broadcast_to(mu, output_shape)
-        param = np.broadcast_to(param, output_shape + param.shape[-1:])
-        if mu.shape[-1] != param.shape[-1]:
-            raise ValueError(f"Shapes for mu and {self._cov_type} don't match")
+        dist_shape = to_tuple(self.shape)
+        mu = broadcast_dist_samples_to(to_shape=dist_shape, samples=[mu], size=size)[0]
+        param = broadcast_dist_samples_to(
+            to_shape=dist_shape + dist_shape[-1:], samples=[param], size=size
+        )[0]
 
         if self._cov_type == "cov":
             chol = np.linalg.cholesky(param)
         elif self._cov_type == "chol":
             chol = param
-        else:
-            inverse = np.linalg.inv(param)
-            chol = np.linalg.cholesky(inverse)
+        else:  # tau -> chol -> swapaxes (chol, -1, -2) -> inv ...
+            lower_chol = np.linalg.cholesky(param)
+            upper_chol = np.swapaxes(lower_chol, -1, -2)
+            chol = np.linalg.inv(upper_chol)
 
+        output_shape = size + dist_shape
         standard_normal = np.random.standard_normal(output_shape)
         return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
 
@@ -404,13 +377,13 @@ def random(self, point=None, size=None):
             nu, mu = draw_values([self.nu, self.mu], point=point, size=size)
             if self._cov_type == "cov":
                 (cov,) = draw_values([self.cov], point=point, size=size)
-                dist = MvNormal.dist(mu=np.zeros_like(mu), cov=cov)
+                dist = MvNormal.dist(mu=np.zeros_like(mu), cov=cov, shape=self.shape)
             elif self._cov_type == "tau":
                 (tau,) = draw_values([self.tau], point=point, size=size)
-                dist = MvNormal.dist(mu=np.zeros_like(mu), tau=tau)
+                dist = MvNormal.dist(mu=np.zeros_like(mu), tau=tau, shape=self.shape)
             else:
                 (chol,) = draw_values([self.chol_cov], point=point, size=size)
-                dist = MvNormal.dist(mu=np.zeros_like(mu), chol=chol)
+                dist = MvNormal.dist(mu=np.zeros_like(mu), chol=chol, shape=self.shape)
 
             samples = dist.random(point, size)
 
@@ -1920,6 +1893,7 @@ def random(self, point=None, size=None):
         """
         # Expand params into terms MvNormal can understand to force consistency
         self._setup_random()
+        self.mv_params["shape"] = self.shape
         dist = MvNormal.dist(**self.mv_params)
         return dist.random(point, size)
 

From ecc9e494c2a44d90be5393717f556b7a3103eafc Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Fri, 27 Nov 2020 20:15:23 +0530
Subject: [PATCH 10/13] Make pylint pass

---
 pymc3/distributions/multivariate.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index a998499cad..3c1a0b7394 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -31,14 +31,7 @@
 
 from pymc3.theanof import floatX
 from . import transforms
-from .distribution import (
-    Continuous,
-    Discrete,
-    draw_values,
-    generate_samples,
-    _DrawValuesContext,
-    is_fast_drawable,
-)
+from .distribution import Continuous, Discrete, draw_values, generate_samples, _DrawValuesContext
 from ..model import Deterministic
 from .continuous import ChiSquared, Normal
 from .special import gammaln, multigammaln

From 3c196880aec6b3e37230fb7317244fec64b764bf Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Fri, 27 Nov 2020 20:59:45 +0530
Subject: [PATCH 11/13] Make test passes :crossed_fingers: hopefully

---
 pymc3/tests/test_mixture.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/tests/test_mixture.py b/pymc3/tests/test_mixture.py
index 27914b6d74..9f016ad333 100644
--- a/pymc3/tests/test_mixture.py
+++ b/pymc3/tests/test_mixture.py
@@ -368,7 +368,7 @@ def build_toy_dataset(N, K):
                     )
                 )
                 chol.append(pm.expand_packed_triangular(D, packed_chol[i], lower=True))
-                comp_dist.append(pm.MvNormal.dist(mu=mu[i], chol=chol[i]))
+                comp_dist.append(pm.MvNormal.dist(mu=mu[i], chol=chol[i], shape=D))
 
             pm.Mixture("x_obs", pi, comp_dist, observed=X)
         with model:

From 6a37514a84255be7f859e369a89106616da196b3 Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Sun, 29 Nov 2020 13:40:39 +0530
Subject: [PATCH 12/13] Modified logic and added tests

---
 pymc3/distributions/multivariate.py      |  17 +++-
 pymc3/tests/test_distributions_random.py | 111 +++++++++++++++++++++++
 2 files changed, 123 insertions(+), 5 deletions(-)

diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 3c1a0b7394..2d4c433d28 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -256,10 +256,18 @@ def random(self, point=None, size=None):
         mu, param = draw_values([self.mu, param_attribute], point=point, size=size)
 
         dist_shape = to_tuple(self.shape)
-        mu = broadcast_dist_samples_to(to_shape=dist_shape, samples=[mu], size=size)[0]
-        param = broadcast_dist_samples_to(
-            to_shape=dist_shape + dist_shape[-1:], samples=[param], size=size
-        )[0]
+        output_shape = size + dist_shape
+
+        # Simple, there can be only be 1 batch dimension, only available from `mu`.
+        # Insert it into `param` before events, if there is a sample shape in front.
+        if param.ndim > 2 and dist_shape[:-1]:
+            param = param.reshape(size + (1,) + param.shape[-2:])
+
+        mu = broadcast_dist_samples_to(to_shape=output_shape, samples=[mu], size=size)[0]
+        param = np.broadcast_to(param, shape=output_shape + dist_shape[-1:])
+
+        assert mu.shape == output_shape
+        assert param.shape == output_shape + dist_shape[-1:]
 
         if self._cov_type == "cov":
             chol = np.linalg.cholesky(param)
@@ -270,7 +278,6 @@ def random(self, point=None, size=None):
             upper_chol = np.swapaxes(lower_chol, -1, -2)
             chol = np.linalg.inv(upper_chol)
 
-        output_shape = size + dist_shape
         standard_normal = np.random.standard_normal(output_shape)
         return mu + np.einsum("...ij,...j->...i", chol, standard_normal)
 
diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py
index b15eb9b00e..4daa1857b8 100644
--- a/pymc3/tests/test_distributions_random.py
+++ b/pymc3/tests/test_distributions_random.py
@@ -12,6 +12,7 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import itertools
 import pytest
 import numpy as np
 import numpy.testing as npt
@@ -27,6 +28,7 @@
     draw_values,
     _DrawValuesContext,
     _DrawValuesContextBlocker,
+    to_tuple,
 )
 from .helpers import SeededTest
 from .test_distributions import (
@@ -1544,3 +1546,112 @@ def test_Triangular(
             prior_samples=prior_samples,
         )
         assert prior["target"].shape == (prior_samples,) + shape
+
+
+def generate_shapes(include_params=False, xfail=False):
+    # fmt: off
+    mudim_as_event = [
+        [None, 1, 3, 10, (10, 3), 100],
+        [(3,)],
+        [(1,), (3,)],
+        ["cov", "chol", "tau"]
+    ]
+    # fmt: on
+    mudim_as_dist = [
+        [None, 1, 3, 10, (10, 3), 100],
+        [(10, 3)],
+        [(1,), (3,), (1, 1), (1, 3), (10, 1), (10, 3)],
+        ["cov", "chol", "tau"],
+    ]
+    if not include_params:
+        del mudim_as_event[-1]
+        del mudim_as_dist[-1]
+    data = itertools.chain(itertools.product(*mudim_as_event), itertools.product(*mudim_as_dist))
+    if xfail:
+        data = list(data)
+        for index in range(len(data)):
+            if data[index][0] in (None, 1):
+                data[index] = pytest.param(
+                    *data[index], marks=pytest.mark.xfail(reason="wait for PR #4214")
+                )
+    return data
+
+
+class TestMvNormal(SeededTest):
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape", "param"],
+        generate_shapes(include_params=True, xfail=False),
+        ids=str,
+    )
+    def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param):
+        dist = pm.MvNormal.dist(mu=np.ones(mu_shape), **{param: np.eye(3)}, shape=dist_shape)
+        output_shape = to_tuple(sample_shape) + dist_shape
+        assert dist.random(size=sample_shape).shape == output_shape
+
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape"],
+        generate_shapes(include_params=False, xfail=True),
+        ids=str,
+    )
+    def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape):
+        with pm.Model() as model:
+            mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape)
+            sd_dist = pm.Exponential.dist(1.0, shape=3)
+            chol, corr, stds = pm.LKJCholeskyCov(
+                "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True
+            )
+            mv = pm.MvNormal("mv", mu, chol=chol, shape=dist_shape)
+            prior = pm.sample_prior_predictive(samples=sample_shape)
+
+        assert prior["mv"].shape == to_tuple(sample_shape) + dist_shape
+
+    @pytest.mark.parametrize(
+        ["sample_shape", "dist_shape", "mu_shape"],
+        generate_shapes(include_params=False, xfail=True),
+        ids=str,
+    )
+    def test_with_cov_rv(self, sample_shape, dist_shape, mu_shape):
+        with pm.Model() as model:
+            mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape)
+            sd_dist = pm.Exponential.dist(1.0, shape=3)
+            chol, corr, stds = pm.LKJCholeskyCov(
+                "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True
+            )
+            mv = pm.MvNormal("mv", mu, cov=pm.math.dot(chol, chol.T), shape=dist_shape)
+            prior = pm.sample_prior_predictive(samples=sample_shape)
+
+        assert prior["mv"].shape == to_tuple(sample_shape) + dist_shape
+
+    def test_issue_3758(self):
+        np.random.seed(42)
+        ndim = 50
+        with pm.Model() as model:
+            a = pm.Normal("a", sigma=100, shape=ndim)
+            b = pm.Normal("b", mu=a, sigma=1, shape=ndim)
+            c = pm.MvNormal("c", mu=a, chol=np.linalg.cholesky(np.eye(ndim)), shape=ndim)
+            d = pm.MvNormal("d", mu=a, cov=np.eye(ndim), shape=ndim)
+            samples = pm.sample_prior_predictive(1000)
+
+        for var in "abcd":
+            assert not np.isnan(np.std(samples[var]))
+
+    def test_issue_3829(self):
+        with pm.Model() as model:
+            x = pm.MvNormal("x", mu=np.zeros(5), cov=np.eye(5), shape=(2, 5))
+            trace_pp = pm.sample_prior_predictive(50)
+
+        assert np.shape(trace_pp["x"][0]) == (2, 5)
+
+    def test_issue_3706(self):
+        N = 10
+        Sigma = np.eye(2)
+
+        with pm.Model() as model:
+
+            X = pm.MvNormal("X", mu=np.zeros(2), cov=Sigma, shape=(N, 2))
+            betas = pm.Normal("betas", 0, 1, shape=2)
+            y = pm.Deterministic("y", pm.math.dot(X, betas))
+
+            prior_pred = pm.sample_prior_predictive(1)
+
+        assert prior_pred["X"].shape == (1, N, 2)

From 3d2ebd7daea63ee9d72e925f6b2878273fc596dd Mon Sep 17 00:00:00 2001
From: Sayam753 <sayamkumar753@yahoo.in>
Date: Sun, 29 Nov 2020 18:31:33 +0530
Subject: [PATCH 13/13] Given a mention in release notes

---
 RELEASE-NOTES.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 7001cb6b4a..4573865189 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -45,6 +45,7 @@ This new version of `Theano-PyMC` comes with an experimental JAX backend which,
 - Enabled the `Multinomial` distribution to handle batch sizes that have more than 2 dimensions. [#4169](https://github.com/pymc-devs/pymc3/pull/4169)
 - Test model logp before starting any MCMC chains (see [#4116](https://github.com/pymc-devs/pymc3/issues/4116))
 - Fix bug in `model.check_test_point` that caused the `test_point` argument to be ignored. (see [PR #4211](https://github.com/pymc-devs/pymc3/pull/4211#issuecomment-727142721))
+- Refactored MvNormal.random method with better handling of sample, batch and event shapes. [#4207](https://github.com/pymc-devs/pymc3/pull/4207)
 
 ### Documentation
 - Added a new notebook demonstrating how to incorporate sampling from a conjugate Dirichlet-multinomial posterior density in conjunction with other step methods (see [#4199](https://github.com/pymc-devs/pymc3/pull/4199)).