Make Metropolis, Slice, PGBART, MetropolisMLDA use point values

Author: brandonwillard

pymc3/aesaraf.py

@@ -11,6 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+from typing import Dict, List
 
 import aesara
 import numpy as np
@@ -222,7 +223,7 @@ def __hash__(self):
         return hash(type(self))
 
 
-def make_shared_replacements(vars, model):
+def make_shared_replacements(point, vars, model):
     """
     Makes shared replacements for all *other* variables than the ones passed.
 
@@ -231,6 +232,7 @@ def make_shared_replacements(vars, model):
 
     Parameters
     ----------
+    point: dictionary mapping variable names to sample values
     vars: list of variables not to make shared
     model: model
 
@@ -240,19 +242,24 @@ def make_shared_replacements(vars, model):
     """
     othervars = set(model.vars) - set(vars)
     return {
-        var: aesara.shared(
-            var.tag.test_value, var.name + "_shared", broadcastable=var.broadcastable
-        )
+        var: aesara.shared(point[var.name], var.name + "_shared", broadcastable=var.broadcastable)
         for var in othervars
     }
 
 
-def join_nonshared_inputs(xs, vars, shared, make_shared=False):
+def join_nonshared_inputs(
+    point: Dict[str, np.ndarray],
+    xs: List[TensorVariable],
+    vars: List[TensorVariable],
+    shared,
+    make_shared: bool = False,
+):
     """
     Takes a list of aesara Variables and joins their non shared inputs into a single input.
 
     Parameters
     ----------
+    point: a sample point
     xs: list of aesara tensors
     vars: list of variables to join
 
@@ -271,17 +278,20 @@ def join_nonshared_inputs(xs, vars, shared, make_shared=False):
         tensor_type = joined.type
         inarray = tensor_type("inarray")
     else:
-        inarray = aesara.shared(joined.tag.test_value, "inarray")
+        if point is None:
+            raise ValueError("A point is required when `make_shared` is True")
+        joined_values = np.concatenate([point[var.name].ravel() for var in vars])
+        inarray = aesara.shared(joined_values, "inarray")
 
-    inarray.tag.test_value = joined.tag.test_value
+    if aesara.config.compute_test_value != "off":
+        inarray.tag.test_value = joined.tag.test_value
 
     replace = {}
     last_idx = 0
     for var in vars:
-        arr_len = at.prod(var.shape)
-        replace[var] = reshape_t(inarray[last_idx : last_idx + arr_len], var.shape).astype(
-            var.dtype
-        )
+        shape = point[var.name].shape
+        arr_len = np.prod(shape, dtype=int)
+        replace[var] = reshape_t(inarray[last_idx : last_idx + arr_len], shape).astype(var.dtype)
         last_idx += arr_len
 
     replace.update(shared)

pymc3/smc/smc.py

@@ -108,12 +108,15 @@ def initialize_population(self):
 
     def setup_kernel(self):
         """Set up the likelihood logp function based on the chosen kernel."""
-        shared = make_shared_replacements(self.variables, self.model)
+        initial_values = self.model.test_point
+        shared = make_shared_replacements(initial_values, self.variables, self.model)
 
         if self.kernel == "abc":
             factors = [var.logpt for var in self.model.free_RVs]
             factors += [at.sum(factor) for factor in self.model.potentials]
-            self.prior_logp_func = logp_forw([at.sum(factors)], self.variables, shared)
+            self.prior_logp_func = logp_forw(
+                initial_values, [at.sum(factors)], self.variables, shared
+            )
             simulator = self.model.observed_RVs[0]
             distance = simulator.distribution.distance
             sum_stat = simulator.distribution.sum_stat
@@ -132,8 +135,12 @@ def setup_kernel(self):
                 self.save_log_pseudolikelihood,
             )
         elif self.kernel == "metropolis":
-            self.prior_logp_func = logp_forw([self.model.varlogpt], self.variables, shared)
-            self.likelihood_logp_func = logp_forw([self.model.datalogpt], self.variables, shared)
+            self.prior_logp_func = logp_forw(
+                initial_values, [self.model.varlogpt], self.variables, shared
+            )
+            self.likelihood_logp_func = logp_forw(
+                initial_values, [self.model.datalogpt], self.variables, shared
+            )
 
     def initialize_logp(self):
         """Initialize the prior and likelihood log probabilities."""
@@ -271,7 +278,7 @@ def posterior_to_trace(self):
         return strace
 
 
-def logp_forw(out_vars, vars, shared):
+def logp_forw(point, out_vars, vars, shared):
     """Compile Aesara function of the model and the input and output variables.
 
     Parameters
@@ -283,7 +290,7 @@ def logp_forw(out_vars, vars, shared):
     shared: List
         containing :class:`aesara.tensor.Tensor` for depended shared data
     """
-    out_list, inarray0 = join_nonshared_inputs(out_vars, vars, shared)
+    out_list, inarray0 = join_nonshared_inputs(point, out_vars, vars, shared)
     f = aesara_function([inarray0], out_list[0])
     f.trust_input = True
     return f

pymc3/step_methods/arraystep.py

@@ -145,12 +145,24 @@ def step(self, point: Dict[str, np.ndarray]):
         if self.allvars:
             inputs.append(point)
 
+        apoint = DictToArrayBijection.map(point)
+        step_res = self.astep(apoint, *inputs)
+
         if self.generates_stats:
-            apoint, stats = self.astep(DictToArrayBijection.map(point), *inputs)
-            return DictToArrayBijection.rmap(apoint), stats
+            apoint_new, stats = step_res
         else:
-            apoint = self.astep(DictToArrayBijection.map(point), *inputs)
-            return DictToArrayBijection.rmap(apoint)
+            apoint_new = step_res
+
+        if not isinstance(apoint_new, RaveledVars):
+            # We assume that the mapping has stayed the same
+            apoint_new = RaveledVars(apoint_new, apoint.point_map_info)
+
+        point_new = DictToArrayBijection.rmap(apoint_new)
+
+        if self.generates_stats:
+            return point_new, stats
+
+        return point_new
 
     def astep(self, apoint, point):
         raise NotImplementedError()
@@ -177,19 +189,41 @@ def __init__(self, vars, shared, blocked=True):
         self.blocked = blocked
 
     def step(self, point):
-        for var, share in self.shared.items():
-            share.set_value(point[var])
+
+        # Remove shared variables from the sample point
+        point_no_shared = point.copy()
+        for name, shared_var in self.shared.items():
+            shared_var.set_value(point[name])
+            if name in point_no_shared:
+                del point_no_shared[name]
+
+        q = DictToArrayBijection.map(point_no_shared)
+
+        step_res = self.astep(q)
 
         if self.generates_stats:
-            apoint, stats = self.astep(DictToArrayBijection.map(point))
-            return DictToArrayBijection.rmap(apoint), stats
+            apoint, stats = step_res
         else:
-            array = DictToArrayBijection.map(point)
-            apoint = self.astep(array)
-            if not isinstance(apoint, RaveledVars):
-                # We assume that the mapping has stayed the same
-                apoint = RaveledVars(apoint, array.point_map_info)
-            return DictToArrayBijection.rmap(apoint)
+            apoint = step_res
+
+        if not isinstance(apoint, RaveledVars):
+            # We assume that the mapping has stayed the same
+            apoint = RaveledVars(apoint, q.point_map_info)
+
+        # We need to re-add the shared variables to the new sample point
+        a_point = DictToArrayBijection.rmap(apoint)
+        new_point = {}
+        for name in point.keys():
+            shared_value = self.shared.get(name, None)
+            if shared_value is not None:
+                new_point[name] = shared_value.get_value()
+            else:
+                new_point[name] = a_point[name]
+
+        if self.generates_stats:
+            return new_point, stats
+
+        return new_point
 
     def astep(self, apoint):
         raise NotImplementedError()