SMC: Fix proposal_dist shape in MH kernel

Author: ricardoV94

pymc3/smc/smc.py

@@ -34,6 +34,7 @@
 from pymc3.blocking import DictToArrayBijection
 from pymc3.model import Point, modelcontext
 from pymc3.sampling import sample_prior_predictive
+from pymc3.step_methods.metropolis import MultivariateNormalProposal
 from pymc3.vartypes import discrete_types
 
 
@@ -449,10 +450,7 @@ def setup_kernel(self):
         Dimension specific scaling is provided by self.proposal_scales and set in self.tune()
         """
         ndim = self.tempered_posterior.shape[1]
-        self.proposal_dist = multivariate_normal(
-            mean=np.zeros(ndim),
-            cov=np.eye(ndim),
-        )
+        self.proposal_dist = MultivariateNormalProposal(np.eye(ndim))
         self.proposal_scales = np.full(self.draws, min(1, 2.38 ** 2 / ndim))
 
     def resample(self):
@@ -477,7 +475,7 @@ def mutate(self):
         for n_step in range(self.n_steps):
             proposal = floatX(
                 self.tempered_posterior
-                + self.proposal_dist.rvs(size=self.draws) * self.proposal_scales[:, None]
+                + self.proposal_dist(num_draws=self.draws) * self.proposal_scales[:, None]
             )
             ll = np.array([self.likelihood_logp_func(prop) for prop in proposal])
             pl = np.array([self.prior_logp_func(prop) for prop in proposal])

pymc3/tests/test_smc.py

@@ -343,3 +343,14 @@ def test_normal_model(self):
         post = idata.posterior.stack(sample=("chain", "draw"))
         assert np.abs(post["mu"].mean() - 10) < 0.1
         assert np.abs(post["sigma"].mean() - 0.5) < 0.05
+
+    def test_proposal_dist_shape(self):
+        with pm.Model() as m:
+            x = pm.Normal("x", 0, 1)
+            y = pm.Normal("y", x, 1, observed=0)
+            trace = pm.sample_smc(
+                draws=10,
+                chains=1,
+                kernel=pm.smc.MH,
+                return_inferencedata=False,
+            )