Make metropolis elemwise updates independent of each other

Not updating q0 after each elemwise update rendered subsequent proposals dependent on the previous ones.

Author: ricardoV94

pymc/step_methods/metropolis.py

@@ -257,14 +257,15 @@ def astep(self, q0: RaveledVars) -> Tuple[RaveledVars, StatsType]:
             q = floatX(q0d + delta)
 
         if self.elemwise_update:
+            q0d = q0d.copy()
             q_temp = q0d.copy()
             # Shuffle order of updates (probably we don't need to do this in every step)
             np.random.shuffle(self.enum_dims)
             for i in self.enum_dims:
                 q_temp[i] = q[i]
                 accept_rate_i = self.delta_logp(q_temp, q0d)
                 q_temp_, accepted_i = metrop_select(accept_rate_i, q_temp, q0d)
-                q_temp[i] = q_temp_[i]
+                q_temp[i] = q0d[i] = q_temp_[i]
                 self.accept_rate_iter[i] = accept_rate_i
                 self.accepted_iter[i] = accepted_i
                 self.accepted_sum[i] += accepted_i

tests/step_methods/test_metropolis.py

@@ -120,6 +120,16 @@ def test_elemwise_update(self, batched_dist):
         assert az.rhat(trace).max()["batched_dist"].values < 1.1
         assert az.ess(trace).min()["batched_dist"].values > 50
 
+    def test_elemwise_update_different_scales(self):
+        mu = [1, 2, 3, 4, 5, 100, 1_000, 10_000]
+        with pm.Model() as m:
+            x = pm.Poisson("x", mu=mu)
+            step = pm.Metropolis([x])
+            trace = pm.sample(draws=1000, chains=2, step=step, random_seed=128).posterior
+
+        np.testing.assert_allclose(trace["x"].mean(("draw", "chain")), mu, rtol=0.1)
+        np.testing.assert_allclose(trace["x"].var(("draw", "chain")), mu, rtol=0.2)
+
     def test_multinomial_no_elemwise_update(self):
         with pm.Model() as m:
             batched_dist = pm.Multinomial("batched_dist", n=5, p=np.ones(4) / 4, shape=(10, 4))