Draft of var_names arg for sample

Author: fonnesbeck

pymc/backends/__init__.py

@@ -67,6 +67,7 @@
 
 import numpy as np
 
+from pytensor.tensor.variable import TensorVariable
 from typing_extensions import TypeAlias
 
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
@@ -99,11 +100,12 @@ def _init_trace(
     stats_dtypes: list[dict[str, type]],
     trace: Optional[BaseTrace],
     model: Model,
+    trace_vars: Optional[list[TensorVariable]] = None,
 ) -> BaseTrace:
     """Initializes a trace backend for a chain."""
     strace: BaseTrace
     if trace is None:
-        strace = NDArray(model=model)
+        strace = NDArray(model=model, vars=trace_vars)
     elif isinstance(trace, BaseTrace):
         if len(trace) > 0:
             raise ValueError("Continuation of traces is no longer supported.")
@@ -123,6 +125,7 @@ def init_traces(
     step: Union[BlockedStep, CompoundStep],
     initial_point: Mapping[str, np.ndarray],
     model: Model,
+    trace_vars: Optional[list[TensorVariable]] = None,
 ) -> tuple[Optional[RunType], Sequence[IBaseTrace]]:
     """Initializes a trace recorder for each chain."""
     if HAS_MCB and isinstance(backend, Backend):
@@ -142,6 +145,7 @@ def init_traces(
             chain_number=chain_number,
             trace=backend,
             model=model,
+            trace_vars=trace_vars,
         )
         for chain_number in range(chains)
     ]

pymc/sampling/mcmc.py

@@ -371,6 +371,7 @@ def sample(
     random_seed: RandomState = None,
     progressbar: bool = True,
     step=None,
+    var_names: Optional[Sequence[str]] = None,
     nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
     initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]] = None,
     init: str = "auto",
@@ -399,6 +400,7 @@ def sample(
     random_seed: RandomState = None,
     progressbar: bool = True,
     step=None,
+    var_names: Optional[Sequence[str]] = None,
     nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
     initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]] = None,
     init: str = "auto",
@@ -427,6 +429,7 @@ def sample(
     random_seed: RandomState = None,
     progressbar: bool = True,
     step=None,
+    var_names: Optional[Sequence[str]] = None,
     nuts_sampler: Literal["pymc", "nutpie", "numpyro", "blackjax"] = "pymc",
     initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]] = None,
     init: str = "auto",
@@ -478,6 +481,8 @@ def sample(
         A step function or collection of functions. If there are variables without step methods,
         step methods for those variables will be assigned automatically. By default the NUTS step
         method will be used, if appropriate to the model.
+    var_names : list of str
+        Names of variables to be monitored. If None, all named variables are selected automatically.
     nuts_sampler : str
         Which NUTS implementation to run. One of ["pymc", "nutpie", "blackjax", "numpyro"].
         This requires the chosen sampler to be installed.
@@ -722,12 +727,18 @@ def sample(
         model.check_start_vals(ip)
         _check_start_shape(model, ip)
 
+    if var_names is not None:
+        trace_vars = [v for v in model.unobserved_RVs if v.name in var_names]
+    else:
+        trace_vars = None
+
     # Create trace backends for each chain
     run, traces = init_traces(
         backend=trace,
         chains=chains,
         expected_length=draws + tune,
         step=step,
+        trace_vars=trace_vars,
         initial_point=ip,
         model=model,
     )
@@ -739,6 +750,7 @@ def sample(
         "traces": traces,
         "chains": chains,
         "tune": tune,
+        "var_names": var_names,
         "progressbar": progressbar,
         "model": model,
         "cores": cores,