feat: Add option not to store some deterministics

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -5,7 +5,7 @@
 from functools import wraps
 from importlib.util import find_spec
 from math import prod
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Literal, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -218,6 +218,7 @@ def make_user_data(shared_vars, shared_data):
 def _compile_pymc_model_numba(
     model: "pm.Model",
     pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    var_names: Iterable[str] | None = None,
     **kwargs,
 ) -> CompiledPyMCModel:
     if find_spec("numba") is None:
@@ -242,6 +243,7 @@ def _compile_pymc_model_numba(
         compute_grad=True,
         join_expanded=True,
         pymc_initial_point_fn=pymc_initial_point_fn,
+        var_names=var_names,
     )
 
     expand_fn = expand_fn_pt.vm.jit_fn
@@ -337,6 +339,7 @@ def _compile_pymc_model_jax(
     *,
     gradient_backend=None,
     pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    var_names: Iterable[str] | None = None,
     **kwargs,
 ):
     if find_spec("jax") is None:
@@ -366,6 +369,7 @@ def _compile_pymc_model_jax(
         compute_grad=gradient_backend == "pytensor",
         join_expanded=False,
         pymc_initial_point_fn=pymc_initial_point_fn,
+        var_names=var_names,
     )
 
     logp_fn = logp_fn_pt.vm.jit_fn
@@ -441,6 +445,7 @@ def compile_pymc_model(
     default_initialization_strategy: Literal[
         "support_point", "prior"
     ] = "support_point",
+    var_names: Iterable[str] | None = None,
     **kwargs,
 ) -> CompiledModel:
     """Compile necessary functions for sampling a pymc model.
@@ -464,6 +469,8 @@ def compile_pymc_model(
     initial_points : dict
         Initial value (strategies) to use instead of what's specified in
         `Model.initial_values`.
+    var_names : list[str] | None
+        A list of variables to store in the trace. If None, store all variables.
     Returns
     -------
     compiled_model : CompiledPyMCModel
@@ -493,13 +500,14 @@ def compile_pymc_model(
         if gradient_backend == "jax":
             raise ValueError("Gradient backend cannot be jax when using numba backend")
         return _compile_pymc_model_numba(
-            model=model, pymc_initial_point_fn=initial_point_fn, **kwargs
+            model=model, pymc_initial_point_fn=initial_point_fn, var_names=var_names, **kwargs
         )
     elif backend.lower() == "jax":
         return _compile_pymc_model_jax(
             model=model,
             gradient_backend=gradient_backend,
             pymc_initial_point_fn=initial_point_fn,
+            var_names=var_names,
             **kwargs,
         )
     else:
@@ -542,6 +550,7 @@ def _make_functions(
     compute_grad: bool,
     join_expanded: bool,
     pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]],
+    var_names: Iterable[str] | None = None,
 ) -> tuple[
     int,
     int,
@@ -568,6 +577,8 @@ def _make_functions(
     pymc_initial_point_fn: Callable
         Initial point function created by
         pymc.initial_point.make_initial_point_fn
+    var_names:
+        Names of variables to store in the trace. Defaults to all variables.
 
     Returns
     -------
@@ -673,6 +684,10 @@ def _make_functions(
         var for var in model.unobserved_value_vars if var.name not in joined_names
     ]
 
+    if var_names is not None:
+        names = set(var_names)
+        remaining_rvs = [var for var in remaining_rvs if var.name in names]
+
     all_names = joined_names + remaining_rvs
 
     all_names = joined_names.copy()

tests/test_pymc.py

@@ -193,6 +193,44 @@ def test_pymc_model_shared(backend, gradient_backend):
         nutpie.sample(compiled3, chains=1)
 
 
+@parameterize_backends
+def test_pymc_var_names(backend, gradient_backend):
+    with pm.Model() as model:
+        mu = pm.Data("mu", -0.1)
+        sigma = pm.Data("sigma", np.ones(3))
+        a = pm.Normal("a", mu=mu, sigma=sigma, shape=3)
+
+        b = pm.Deterministic("b", mu * a)
+        pm.Deterministic("c", mu * b)
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend, var_names=None,
+    )
+    trace = nutpie.sample(compiled, chains=1, seed=1)
+
+    # Check that variables are stored
+    assert hasattr(trace.posterior, "b")
+    assert hasattr(trace.posterior, "c")
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend, var_names=[],
+    )
+    trace = nutpie.sample(compiled, chains=1, seed=1)
+
+    # Check that variables are stored
+    assert not hasattr(trace.posterior, "b")
+    assert not hasattr(trace.posterior, "c")
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend, var_names=["b"],
+    )
+    trace = nutpie.sample(compiled, chains=1, seed=1)
+
+    # Check that variables are stored
+    assert hasattr(trace.posterior, "b")
+    assert not hasattr(trace.posterior, "c")
+
+
 @pytest.mark.parametrize(
     ("backend", "gradient_backend"),
     [