Implement get_value_vars_from_user_vars util to check sampler vars are valid

Deprecates old `allinmodel`

Author: ricardoV94

pymc/model.py

@@ -69,6 +69,7 @@
     UNSET,
     WithMemoization,
     get_transformed_name,
+    get_value_vars_from_user_vars,
     get_var_name,
     treedict,
     treelist,
@@ -617,6 +618,7 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs):
         if grad_vars is None:
             grad_vars = self.continuous_value_vars
         else:
+            grad_vars = get_value_vars_from_user_vars(grad_vars, self)
             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}")

pymc/step_methods/hmc/base_hmc.py

@@ -31,6 +31,7 @@
 from pymc.step_methods.hmc import integration
 from pymc.step_methods.hmc.quadpotential import QuadPotentialDiagAdapt, quad_potential
 from pymc.tuning import guess_scaling
+from pymc.util import get_value_vars_from_user_vars
 
 logger = logging.getLogger("pymc")
 
@@ -91,8 +92,7 @@ def __init__(
         if vars is None:
             vars = self._model.continuous_value_vars
         else:
-            vars = [self._model.rvs_to_values.get(var, var) for var in vars]
-
+            vars = get_value_vars_from_user_vars(vars, self._model)
         super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **aesara_kwargs)
 
         self.adapt_step_size = adapt_step_size

pymc/step_methods/metropolis.py

@@ -56,6 +56,8 @@
     "MultivariateNormalProposal",
 ]
 
+from pymc.util import get_value_vars_from_user_vars
+
 # Available proposal distributions for Metropolis
 
 
@@ -176,9 +178,7 @@ 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)
+            vars = get_value_vars_from_user_vars(vars, model)
 
         initial_values_shape = [initial_values[v.name].shape for v in vars]
         if S is None:
@@ -394,7 +394,7 @@ 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]
+        vars = get_value_vars_from_user_vars(vars, model)
 
         if not all([v.dtype in pm.discrete_types for v in vars]):
             raise ValueError("All variables must be Bernoulli for BinaryMetropolis")
@@ -484,8 +484,9 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
         # transition probabilities
         self.transit_p = transit_p
 
+        vars = get_value_vars_from_user_vars(vars, model)
+
         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":
@@ -566,8 +567,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)
+        vars = get_value_vars_from_user_vars(vars, model)
 
         initial_point = model.initial_point()
 
@@ -777,8 +777,7 @@ def __init__(
         if vars is None:
             vars = model.continuous_value_vars
         else:
-            vars = [model.rvs_to_values.get(var, var) for var in vars]
-        vars = pm.inputvars(vars)
+            vars = get_value_vars_from_user_vars(vars, model)
 
         if S is None:
             S = np.ones(initial_values_size)
@@ -928,8 +927,7 @@ def __init__(
         if vars is None:
             vars = model.continuous_value_vars
         else:
-            vars = [model.rvs_to_values.get(var, var) for var in vars]
-        vars = pm.inputvars(vars)
+            vars = get_value_vars_from_user_vars(vars, model)
 
         if S is None:
             S = np.ones(initial_values_size)

pymc/step_methods/slicer.py

@@ -17,10 +17,10 @@
 import numpy as np
 import numpy.random as nr
 
-from pymc.aesaraf import inputvars
 from pymc.blocking import RaveledVars
 from pymc.model import modelcontext
 from pymc.step_methods.arraystep import ArrayStep, Competence
+from pymc.util import get_value_vars_from_user_vars
 from pymc.vartypes import continuous_types
 
 __all__ = ["Slice"]
@@ -65,8 +65,7 @@ def __init__(self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, *
         if vars is None:
             vars = self.model.continuous_value_vars
         else:
-            vars = [self.model.rvs_to_values.get(var, var) for var in vars]
-        vars = inputvars(vars)
+            vars = get_value_vars_from_user_vars(vars, self.model)
 
         super().__init__(vars, [self.model.compile_logp()], **kwargs)
 

pymc/tests/test_util.py

@@ -28,6 +28,7 @@
     _get_seeds_per_chain,
     dataset_to_point_list,
     drop_warning_stat,
+    get_value_vars_from_user_vars,
     hash_key,
     hashable,
     locally_cachedmethod,
@@ -236,3 +237,32 @@ def test_get_seeds_per_chain():
 
     with pytest.raises(ValueError, match=re.escape("The `seeds` must be array-like")):
         _get_seeds_per_chain({1: 1, 2: 2}, 2)
+
+
+def test_get_value_vars_from_user_vars():
+    with pm.Model() as model1:
+        x1 = pm.Normal("x1", mu=0, sigma=1)
+        y1 = pm.Normal("y1", mu=0, sigma=1)
+
+    x1_value = model1.rvs_to_values[x1]
+    y1_value = model1.rvs_to_values[y1]
+    assert get_value_vars_from_user_vars([x1, y1], model1) == [x1_value, y1_value]
+    assert get_value_vars_from_user_vars([x1], model1) == [x1_value]
+    # The next line does not wrap the variable in a list on purpose, to test the
+    # utility function can handle those as promised
+    assert get_value_vars_from_user_vars(x1_value, model1) == [x1_value]
+
+    with pm.Model() as model2:
+        x2 = pm.Normal("x2", mu=0, sigma=1)
+        y2 = pm.Normal("y2", mu=0, sigma=1)
+        det2 = pm.Deterministic("det2", x2 + y2)
+
+    prefix = "The following variables are not random variables in the model:"
+    with pytest.raises(ValueError, match=rf"{prefix} \['x2', 'y2'\]"):
+        get_value_vars_from_user_vars([x2, y2], model1)
+    with pytest.raises(ValueError, match=rf"{prefix} \['x2'\]"):
+        get_value_vars_from_user_vars([x2, y1], model1)
+    with pytest.raises(ValueError, match=rf"{prefix} \['x2'\]"):
+        get_value_vars_from_user_vars([x2], model1)
+    with pytest.raises(ValueError, match=rf"{prefix} \['det2'\]"):
+        get_value_vars_from_user_vars([det2], model2)

pymc/tests/tuning/test_starting.py

@@ -22,7 +22,7 @@
 from pymc.tests.checks import close_to
 from pymc.tests.helpers import select_by_precision
 from pymc.tests.models import non_normal, simple_arbitrary_det, simple_model
-from pymc.tuning import find_MAP, starting
+from pymc.tuning import find_MAP
 
 
 @pytest.mark.parametrize("bounded", [False, True])
@@ -147,28 +147,3 @@ def test_find_MAP_issue_4488():
     assert not set.difference({"x_missing", "x_missing_log__", "y"}, set(map_estimate.keys()))
     np.testing.assert_allclose(map_estimate["x_missing"], 0.2, rtol=1e-4, atol=1e-4)
     np.testing.assert_allclose(map_estimate["y"], [2.0, map_estimate["x_missing"][0] + 1])
-
-
-def test_allinmodel():
-    model1 = pm.Model()
-    model2 = pm.Model()
-    with model1:
-        x1 = pm.Normal("x1", mu=0, sigma=1)
-        y1 = pm.Normal("y1", mu=0, sigma=1)
-    with model2:
-        x2 = pm.Normal("x2", mu=0, sigma=1)
-        y2 = pm.Normal("y2", mu=0, sigma=1)
-
-    x1 = model1.rvs_to_values[x1]
-    y1 = model1.rvs_to_values[y1]
-    x2 = model2.rvs_to_values[x2]
-    y2 = model2.rvs_to_values[y2]
-
-    starting.allinmodel([x1, y1], model1)
-    starting.allinmodel([x1], model1)
-    with pytest.raises(ValueError, match=r"Some variables not in the model: \['x2', 'y2'\]"):
-        starting.allinmodel([x2, y2], model1)
-    with pytest.raises(ValueError, match=r"Some variables not in the model: \['x2'\]"):
-        starting.allinmodel([x2, y1], model1)
-    with pytest.raises(ValueError, match=r"Some variables not in the model: \['x2'\]"):
-        starting.allinmodel([x2], model1)

pymc/tuning/starting.py

@@ -30,11 +30,10 @@
 
 import pymc as pm
 
-from pymc.aesaraf import inputvars
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
-from pymc.util import get_default_varnames, get_var_name
+from pymc.util import get_default_varnames, get_value_vars_from_user_vars
 from pymc.vartypes import discrete_types, typefilter
 
 __all__ = ["find_MAP"]
@@ -96,11 +95,9 @@ def find_MAP(
         if not vars:
             raise ValueError("Model has no unobserved continuous variables.")
     else:
-        vars = [model.rvs_to_values.get(var, var) for var in vars]
+        vars = get_value_vars_from_user_vars(vars, model)
 
-    vars = inputvars(vars)
     disc_vars = list(typefilter(vars, discrete_types))
-    allinmodel(vars, model)
     ipfn = make_initial_point_fn(
         model=model,
         jitter_rvs=set(),
@@ -182,13 +179,6 @@ def allfinite(x):
     return np.all(isfinite(x))
 
 
-def allinmodel(vars, model):
-    notin = [v for v in vars if v not in model.value_vars]
-    if notin:
-        notin = list(map(get_var_name, notin))
-        raise ValueError("Some variables not in the model: " + str(notin))
-
-
 class CostFuncWrapper:
     def __init__(self, maxeval=5000, progressbar=True, logp_func=None, dlogp_func=None):
         self.n_eval = 0

pymc/util.py

@@ -21,6 +21,7 @@
 import numpy as np
 import xarray
 
+from aesara import Variable
 from aesara.compile import SharedVariable
 from cachetools import LRUCache, cachedmethod
 
@@ -441,3 +442,42 @@ def _get_unique_seeds_per_chain(integers_fn):
         )
 
     return random_state
+
+
+def get_value_vars_from_user_vars(
+    vars: Union[Variable, Sequence[Variable]], model
+) -> List[Variable]:
+    """This function converts user "vars" input into value variables
+
+    More often than not, users will pass random variables, and we will extract the
+    respective value variables, but we also allow for the input to already be value
+    variables, in case the function is called internally or by a "super-user"
+
+    Returns
+    -------
+    value_vars: list of TensorVariable
+        List of model value variables that correspond to the input vars
+
+    Raises
+    ------
+    ValueError:
+        If any of the provided variables do not correspond to any model value variable
+    """
+    if not isinstance(vars, Sequence):
+        # Single var was passed
+        value_vars = [model.rvs_to_values.get(vars, vars)]
+    else:
+        value_vars = [model.rvs_to_values.get(var, var) for var in vars]
+
+    # Check that we only have value vars from the model
+    model_value_vars = model.value_vars
+    notin = [v for v in value_vars if v not in model_value_vars]
+    if notin:
+        notin = list(map(get_var_name, notin))
+        # We mention random variables, even though the input may be a wrong value variable
+        # because most users don't know about that duality
+        raise ValueError(
+            "The following variables are not random variables in the model: " + str(notin)
+        )
+
+    return value_vars