Repository: pymc-devs/pymc
Pull request: Add var_names argument to sample
## Description

Allow for filtering of variables included in sampled trace via an optional `var_names` argument, similar to what is done for plotting.

## Related Issue
<!--- It is good practice to first open an issue explaining the bug / new feature that is addressed by this PR -->
<!--- Please type an `x` in one of the boxes below and provide the issue number after the # sign: -->
- [x] Closes #7068 
- [ ] Related to #

## Checklist
<!--- Make sure you have completed the following steps before submitting your PR -->
<!--- Feel free to type an `x` in all the boxes below to let us know you have completed the steps: -->
- [x] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [x] Included tests that prove the fix is effective or that the new feature works
- [x] Added necessary documentation (docstrings and/or example notebooks)
- [ ] If you are a pro: each commit corresponds to a [relevant logical change](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->

## Type of change
<!--- Select one of the categories below by typing an `x` in the box -->
- [x] New feature / enhancement
- [ ] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):
<!--- Additionally, if you are a maintainer or reviewer, please make sure that the appropriate labels are added to this PR -->


<!-- readthedocs-preview pymc start -->
----
📚 Documentation preview 📚: https://pymc--7206.org.readthedocs.build/en/7206/

<!-- readthedocs-preview pymc end -->

Problem statement:
Allow storing subset of variable in `pm.sample` via a `var_names` kwarg
### Description

Sometimes models with large Deterministics run out of RAM during sampling. One could compute the Deterministics afterwards easily, but would need to create the same model twice, once with the Deterministics and once without.

Hints from issue discussion:
Also a helper `pm.compute_deterministics` to compute any deterministics on demand

Commit message: Added jax test

Patch to review:
diff --git a/pymc/backends/__init__.py b/pymc/backends/__init__.py
index b274af10b6..b63f68acc6 100644
--- a/pymc/backends/__init__.py
+++ b/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)
     ]
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
index 01f8b0d502..03627ea890 100644
--- a/pymc/sampling/jax.py
+++ b/pymc/sampling/jax.py
@@ -532,15 +532,17 @@ def sample_jax_nuts(
 
     model = modelcontext(model)
 
-    if var_names is None:
-        var_names = model.unobserved_value_vars
+    if var_names is not None:
+        vars_ = [v for v in model.unobserved_value_vars if v.name in var_names]
+    else:
+        vars_ = model.unobserved_value_vars
 
     if nuts_kwargs is None:
         nuts_kwargs = {}
     else:
         nuts_kwargs = nuts_kwargs.copy()
 
-    vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
+    vars_to_sample = list(get_default_varnames(vars_, include_transformed=keep_untransformed))
 
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
index 313a6ab8c0..d2c9f051c7 100644
--- a/pymc/sampling/mcmc.py
+++ b/pymc/sampling/mcmc.py
@@ -264,6 +264,7 @@ def _sample_external_nuts(
     random_seed: Union[RandomState, None],
     initvals: Union[StartDict, Sequence[Optional[StartDict]], None],
     model: Model,
+    var_names: Optional[Sequence[str]],
     progressbar: bool,
     idata_kwargs: Optional[dict],
     nuts_sampler_kwargs: Optional[dict],
@@ -348,6 +349,7 @@ def _sample_external_nuts(
             random_seed=random_seed,
             initvals=initvals,
             model=model,
+            var_names=var_names,
             progressbar=progressbar,
             nuts_sampler=sampler,
             idata_kwargs=idata_kwargs,
@@ -371,6 +373,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 +402,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 +431,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 +483,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, optional
+        Names of variables to be stored in the trace. Defaults to all free variables and deterministics.
     nuts_sampler : str
         Which NUTS implementation to run. One of ["pymc", "nutpie", "blackjax", "numpyro"].
         This requires the chosen sampler to be installed.
@@ -680,6 +687,7 @@ def sample(
             random_seed=random_seed,
             initvals=initvals,
             model=model,
+            var_names=var_names,
             progressbar=progressbar,
             idata_kwargs=idata_kwargs,
             nuts_sampler_kwargs=nuts_sampler_kwargs,
@@ -722,12 +730,19 @@ 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]
+        assert len(trace_vars) == len(var_names), "Not all var_names were found in the model"
+    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 +754,7 @@ def sample(
         "traces": traces,
         "chains": chains,
         "tune": tune,
+        "var_names": var_names,
         "progressbar": progressbar,
         "model": model,
         "cores": cores,
diff --git a/tests/sampling/test_jax.py b/tests/sampling/test_jax.py
index d8d0cae246..77121b37ce 100644
--- a/tests/sampling/test_jax.py
+++ b/tests/sampling/test_jax.py
@@ -491,6 +491,15 @@ def test_sample_partially_observed():
     assert idata.posterior["x"].shape == (1, 10, 3)
 
 
+def test_sample_var_names():
+    with pm.Model() as model:
+        a = pm.Normal("a")
+        b = pm.Deterministic("b", a**2)
+        idata = pm.sample(10, tune=10, nuts_sampler="numpyro", var_names=["a"])
+        assert "a" in idata.posterior
+        assert "b" not in idata.posterior
+
+
 @pytest.mark.parametrize("nuts_sampler", ("numpyro", "blackjax"))
 def test_convergence_warnings(caplog, nuts_sampler):
     with pm.Model() as m:
diff --git a/tests/sampling/test_mcmc.py b/tests/sampling/test_mcmc.py
index 0fc03dd631..a18430818d 100644
--- a/tests/sampling/test_mcmc.py
+++ b/tests/sampling/test_mcmc.py
@@ -694,6 +694,15 @@ def test_no_init_nuts_compound(caplog):
         assert "Initializing NUTS" not in caplog.text
 
 
+def test_sample_var_names():
+    with pm.Model() as model:
+        a = pm.Normal("a")
+        b = pm.Deterministic("b", a**2)
+        idata = pm.sample(10, tune=10, var_names=["a"])
+        assert "a" in idata.posterior
+        assert "b" not in idata.posterior
+
+
 class TestAssignStepMethods:
     def test_bernoulli(self):
         """Test bernoulli distribution is assigned binary gibbs metropolis method"""


Code review comments:
diff_hunk: @@ -532,15 +532,17 @@ def sample_jax_nuts(
 
     model = modelcontext(model)
 
-    if var_names is None:
-        var_names = model.unobserved_value_vars
+    if var_names is not None:
+        vars_ = [v for v in model.unobserved_value_vars if v.name in var_names]
path: pymc/sampling/jax.py
text: @user1:
A more readable name? I guess mypy didn't like the redefinition?
```suggestion
        filtered_var_names = [v for v in model.unobserved_value_vars if v.name in var_names]
```