Handle RVs assigned to steps

pymc3/model.py

@@ -714,14 +714,11 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs):
             `alpha` can be changed using `ValueGradFunction.set_weights([alpha])`.
         """
         if grad_vars is None:
-            grad_vars = [v.tag.value_var for v in typefilter(self.free_RVs, continuous_types)]
+            grad_vars = [self.rvs_to_values[v] for v in typefilter(self.free_RVs, continuous_types)]
         else:
             for i, var in enumerate(grad_vars):
                 if var.dtype not in continuous_types:
                     raise ValueError(f"Can only compute the gradient of continuous types: {var}")
-                # We allow one to pass the random variable terms as arguments
-                if hasattr(var.tag, "value_var"):
-                    grad_vars[i] = var.tag.value_var
 
         if tempered:
             with self:

pymc3/sampling.py

@@ -187,11 +187,7 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None
         except TypeError:
             steps.append(step)
         for step in steps:
-            try:
-                assigned_vars = assigned_vars.union(set(step.vars))
-            except AttributeError:
-                for method in step.methods:
-                    assigned_vars = assigned_vars.union(set(method.vars))
+            assigned_vars = assigned_vars.union(set(step.vars))
 
     # Use competence classmethods to select step methods for remaining
     # variables

pymc3/step_methods/__init__.py

@@ -14,7 +14,6 @@
 
 from pymc3.step_methods.compound import CompoundStep
 from pymc3.step_methods.elliptical_slice import EllipticalSlice
-from pymc3.step_methods.gibbs import ElemwiseCategorical
 from pymc3.step_methods.hmc import NUTS, HamiltonianMC
 from pymc3.step_methods.metropolis import (
     BinaryGibbsMetropolis,

pymc3/step_methods/arraystep.py

@@ -134,7 +134,7 @@ class ArrayStep(BlockedStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     fs: list of logp Aesara functions
     allvars: Boolean (default False)
     blocked: Boolean (default True)
@@ -190,7 +190,7 @@ def __init__(self, vars, shared, blocked=True):
         """
         Parameters
         ----------
-        vars: list of sampling variables
+        vars: list of sampling value variables
         shared: dict of Aesara variable -> shared variable
         blocked: Boolean (default True)
         """
@@ -235,7 +235,7 @@ def __init__(self, vars, shared, blocked=True):
         """
         Parameters
         ----------
-        vars: list of sampling variables
+        vars: list of sampling value variables
         shared: dict of Aesara variable -> shared variable
         blocked: Boolean (default True)
         """

pymc3/step_methods/compound.py

@@ -71,3 +71,7 @@ def reset_tuning(self):
         for method in self.methods:
             if hasattr(method, "reset_tuning"):
                 method.reset_tuning()
+
+    @property
+    def vars(self):
+        return [var for method in self.methods for var in method.vars]

pymc3/step_methods/elliptical_slice.py

@@ -16,6 +16,7 @@
 import numpy as np
 import numpy.random as nr
 
+from pymc3.aesaraf import inputvars
 from pymc3.model import modelcontext
 from pymc3.step_methods.arraystep import ArrayStep, Competence
 
@@ -61,7 +62,7 @@ class EllipticalSlice(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     prior_cov: array, optional
         Covariance matrix of the multivariate Gaussian prior.
     prior_chol: array, optional
@@ -88,6 +89,8 @@ def __init__(self, vars=None, prior_cov=None, prior_chol=None, model=None, **kwa
 
         if vars is None:
             vars = self.model.cont_vars
+        else:
+            vars = [self.model.rvs_to_values.get(var, var) for var in vars]
         vars = inputvars(vars)
 
         super().__init__(vars, [self.model.fastlogp], **kwargs)

pymc3/step_methods/hmc/base_hmc.py

@@ -89,6 +89,8 @@ def __init__(
 
         if vars is None:
             vars = self._model.cont_vars
+        else:
+            vars = [self._model.rvs_to_values.get(var, var) for var in vars]
 
         super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **aesara_kwargs)
 

pymc3/step_methods/hmc/hmc.py

@@ -60,7 +60,7 @@ def __init__(self, vars=None, path_length=2.0, max_steps=1024, **kwargs):
         Parameters
         ----------
         vars: list, default=None
-            List of Aesara variables. If None, all continuous RVs from the
+            List of value variables. If None, all continuous RVs from the
             model are included.
         path_length: float, default=2
             Total length to travel

pymc3/step_methods/hmc/nuts.py

@@ -115,7 +115,7 @@ def __init__(self, vars=None, max_treedepth=10, early_max_treedepth=8, **kwargs)
         Parameters
         ----------
         vars: list, default=None
-            List of Aesara variables. If None, all continuous RVs from the
+            List of value variables. If None, all continuous RVs from the
             model are included.
         Emax: float, default 1000
             Maximum energy change allowed during leapfrog steps. Larger

pymc3/step_methods/metropolis.py

@@ -130,7 +130,7 @@ def __init__(
         Parameters
         ----------
         vars: list
-            List of variables for sampler
+            List of value variables for sampler
         S: standard deviation or covariance matrix
             Some measure of variance to parameterize proposal distribution
         proposal_dist: function
@@ -153,6 +153,8 @@ def __init__(
 
         if vars is None:
             vars = model.value_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:
@@ -288,7 +290,7 @@ class BinaryMetropolis(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     scaling: scalar or array
         Initial scale factor for proposal. Defaults to 1.
     tune: bool
@@ -321,6 +323,8 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None):
         self.steps_until_tune = tune_interval
         self.accepted = 0
 
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
+
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError("All variables must be Bernoulli for BinaryMetropolis")
 
@@ -388,7 +392,7 @@ class BinaryGibbsMetropolis(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     order: list or 'random'
         List of integers indicating the Gibbs update order
         e.g., [0, 2, 1, ...]. Default is random
@@ -410,6 +414,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
         self.transit_p = transit_p
 
         initial_point = model.initial_point
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
         self.dim = sum(initial_point[v.name].size for v in vars)
 
         if order == "random":
@@ -490,6 +495,7 @@ def __init__(self, vars, proposal="uniform", order="random", model=None):
 
         model = pm.modelcontext(model)
 
+        vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         initial_point = model.initial_point
@@ -697,6 +703,8 @@ def __init__(
 
         if vars is None:
             vars = model.cont_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:
@@ -846,6 +854,8 @@ def __init__(
 
         if vars is None:
             vars = model.cont_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
         vars = pm.inputvars(vars)
 
         if S is None:

pymc3/step_methods/mlda.py

@@ -74,6 +74,8 @@ def __init__(self, *args, **kwargs):
             value_vars = kwargs.get("vars", None)
             if value_vars is None:
                 value_vars = model.value_vars
+            else:
+                value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
             value_vars = pm.inputvars(value_vars)
             shared = pm.make_shared_replacements(initial_values, value_vars, model)
 
@@ -142,6 +144,8 @@ def __init__(self, *args, **kwargs):
             value_vars = kwargs.get("vars", None)
             if value_vars is None:
                 value_vars = model.value_vars
+            else:
+                value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
             value_vars = pm.inputvars(value_vars)
             shared = pm.make_shared_replacements(initial_values, value_vars, model)
 
@@ -218,7 +222,7 @@ class MLDA(ArrayStepShared):
         Note this list excludes the model passed to the model
         argument above, which is the finest available.
     vars : list
-        List of variables for sampler
+        List of value variables for sampler
     base_sampler : string
         Sampler used in the base (coarsest) chain. Can be 'Metropolis' or
         'DEMetropolisZ'. Defaults to 'DEMetropolisZ'.
@@ -549,6 +553,8 @@ def __init__(
         # Process model variables
         if value_vars is None:
             value_vars = model.value_vars
+        else:
+            value_vars = [model.rvs_to_values.get(var, var) for var in value_vars]
         value_vars = pm.inputvars(value_vars)
         self.vars = value_vars
         self.var_names = [var.name for var in self.vars]

pymc3/step_methods/pgbart.py

@@ -39,7 +39,7 @@ class PGBART(ArrayStepShared):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     num_particles : int
         Number of particles for the conditional SMC sampler. Defaults to 10
     max_stages : int

pymc3/step_methods/sgmcmc.py

@@ -87,7 +87,7 @@ class BaseStochasticGradient(ArrayStepShared):
     Parameters
     ----------
     vars: list
-        List of variables for sampler
+        List of value variables for sampler
     batch_size`: int
         Batch Size for each step
     total_size: int
@@ -132,6 +132,8 @@ def __init__(
 
         if vars is None:
             vars = model.value_vars
+        else:
+            vars = [model.rvs_to_values.get(var, var) for var in vars]
 
         vars = inputvars(vars)
 

pymc3/step_methods/slicer.py

@@ -35,7 +35,7 @@ class Slice(ArrayStep):
     Parameters
     ----------
     vars: list
-        List of variables for sampler.
+        List of value variables for sampler.
     w: float
         Initial width of slice (Defaults to 1).
     tune: bool
@@ -57,6 +57,8 @@ def __init__(self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, *
 
         if vars is None:
             vars = self.model.cont_vars
+        else:
+            vars = [self.model.rvs_to_values.get(var, var) for var in vars]
         vars = inputvars(vars)
 
         super().__init__(vars, [self.model.fastlogp], **kwargs)