get_vars_in_point_list only considers variables that are both in the model and in the point list

Author: lucianopaz

pymc/model.py

@@ -1601,6 +1601,9 @@ def __getitem__(self, key):
             except KeyError:
                 raise e
 
+    def __contains__(self, key):
+        return key in self.named_vars or self.name_for(key) in self.named_vars
+
     def compile_fn(
         self,
         outs: Sequence[Variable],

pymc/sampling.py

@@ -1592,7 +1592,7 @@ def get_vars_in_point_list(trace, model):
         names_in_trace = list(trace[0])
     else:
         names_in_trace = trace.varnames
-    vars_in_trace = [model[v] for v in names_in_trace]
+    vars_in_trace = [model[v] for v in names_in_trace if v in model]
     return vars_in_trace
 
 

pymc/tests/test_sampling.py

@@ -48,6 +48,7 @@
     _get_seeds_per_chain,
     assign_step_methods,
     compile_forward_sampling_function,
+    get_vars_in_point_list,
 )
 from pymc.step_methods import (
     NUTS,
@@ -2628,3 +2629,24 @@ def test_sample(self):
         np.testing.assert_allclose(
             x_pred, pp_trace1.posterior_predictive["obs"].mean(("chain", "draw")), atol=1e-1
         )
+
+
+def test_get_vars_in_point_list():
+    with pm.Model() as modelA:
+        pm.Normal("a", 0, 1)
+        pm.Normal("b", 0, 1)
+    with pm.Model() as modelB:
+        a = pm.Normal("a", 0, 1)
+        pm.Normal("c", 0, 1)
+
+    point_list = [{"a": 0, "b": 0}]
+    vars_in_trace = get_vars_in_point_list(point_list, modelB)
+    assert set(vars_in_trace) == {a}
+
+    strace = pm.backends.NDArray(model=modelB, vars=modelA.free_RVs)
+    strace.setup(1, 1)
+    strace.values = point_list[0]
+    strace.draw_idx = 1
+    trace = MultiTrace([strace])
+    vars_in_trace = get_vars_in_point_list(trace, modelB)
+    assert set(vars_in_trace) == {a}