Keep broadcasting information in make_shared_replacements (pymc-devs#4526)

Author: michaelosthege

RELEASE-NOTES.md

@@ -9,8 +9,8 @@
 + ...
 
 ### Maintenance
-- ⚠ Our memoization mechanism wasn't robust against hash collisions (#4506), sometimes resulting in incorrect values in, for example, posterior predictives. The `pymc3.memoize` module was removed and replaced with `cachetools`.  The `hashable` function and `WithMemoization` class were moved to `pymc3.util`.
-- ...
+- ⚠ Our memoization mechanism wasn't robust against hash collisions (#4506), sometimes resulting in incorrect values in, for example, posterior predictives. The `pymc3.memoize` module was removed and replaced with `cachetools`.  The `hashable` function and `WithMemoization` class were moved to `pymc3.util` (see #4525).
+- `pm.make_shared_replacements` now retains broadcasting information which fixes issues with Metropolis samplers (see [#4492](https://github.com/pymc-devs/pymc3/pull/4492)).
 
 ## PyMC3 3.11.1 (12 February 2021)
 

pymc3/tests/test_theanof.py

@@ -19,13 +19,38 @@
 import theano
 import theano.tensor as tt
 
+import pymc3 as pm
+
 from pymc3.theanof import _conversion_map, take_along_axis
 from pymc3.vartypes import int_types
 
 FLOATX = str(theano.config.floatX)
 INTX = str(_conversion_map[FLOATX])
 
 
+class TestBroadcasting:
+    def test_make_shared_replacements(self):
+        """Check if pm.make_shared_replacements preserves broadcasting."""
+
+        with pm.Model() as test_model:
+            test1 = pm.Normal("test1", mu=0.0, sigma=1.0, shape=(1, 10))
+            test2 = pm.Normal("test2", mu=0.0, sigma=1.0, shape=(10, 1))
+
+        # Replace test1 with a shared variable, keep test 2 the same
+        replacement = pm.make_shared_replacements([test_model.test2], test_model)
+        assert test_model.test1.broadcastable == replacement[test_model.test1].broadcastable
+
+    def test_metropolis_sampling(self):
+        """Check if the Metropolis sampler can handle broadcasting."""
+        with pm.Model() as test_model:
+            test1 = pm.Normal("test1", mu=0.0, sigma=1.0, shape=(1, 10))
+            test2 = pm.Normal("test2", mu=test1, sigma=1.0, shape=(10, 10))
+
+            step = pm.Metropolis()
+            # This should fail immediately if broadcasting does not work.
+            pm.sample(tune=5, draws=7, cores=1, step=step, compute_convergence_checks=False)
+
+
 def _make_along_axis_idx(arr_shape, indices, axis):
     # compute dimensions to iterate over
     if str(indices.dtype) not in int_types:

pymc3/theanof.py

@@ -235,7 +235,12 @@ def make_shared_replacements(vars, model):
     Dict of variable -> new shared variable
     """
     othervars = set(model.vars) - set(vars)
-    return {var: theano.shared(var.tag.test_value, var.name + "_shared") for var in othervars}
+    return {
+        var: theano.shared(
+            var.tag.test_value, var.name + "_shared", broadcastable=var.broadcastable
+        )
+        for var in othervars
+    }
 
 
 def join_nonshared_inputs(xs, vars, shared, make_shared=False):