Keep the remaining caching in FrozenModel

pymc/backends/base.py

@@ -34,7 +34,6 @@
 
 from pymc.backends.report import SamplerReport
 from pymc.model import modelcontext
-from pymc.pytensorf import compile
 from pymc.util import get_var_name
 
 logger = logging.getLogger(__name__)
@@ -183,9 +182,11 @@ def __init__(
 
         if fn is None:
             # borrow=True avoids deepcopy when inputs=output which is the case for untransformed value variables
-            fn = compile(
+            # Routed through the model so the compilation is reused on frozen models.
+            fn = model.compile_fn(
+                outs=[pytensor.Out(v, borrow=True) for v in vars],
                 inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
-                outputs=[pytensor.Out(v, borrow=True) for v in vars],
+                point_fn=False,
                 on_unused_input="ignore",
             )
             fn.trust_input = True

pymc/initial_point.py

@@ -104,17 +104,15 @@ def make_initial_point_fns_per_chain(
         # One strategy for all chains
         # Only one function compilation is needed.
         ipfns = [
-            make_initial_point_fn(
-                model=model,
+            model._initial_point_fn(
                 overrides=overrides,
                 jitter_rvs=jitter_rvs,
                 return_transformed=True,
             )
         ] * chains
     elif len(overrides) == chains:
         ipfns = [
-            make_initial_point_fn(
-                model=model,
+            model._initial_point_fn(
                 jitter_rvs=jitter_rvs,
                 overrides=chain_overrides,
                 return_transformed=True,

pymc/model/core.py

@@ -138,6 +138,33 @@ def modelcontext(model: BaseModel | None) -> BaseModel:
     return model
 
 
+def _make_value_grad_function(
+    model, grad_vars, *, tempered=False, ravel_inputs=None, initial_point, **kwargs
+) -> ValueGradFunction:
+    """Build the logp/dlogp function for ``grad_vars``, treating the rest as extra inputs."""
+    grad_vars = list(grad_vars)
+    if tempered:
+        costs = [model.varlogp, model.datalogp]
+    else:
+        costs = [model.logp()]
+
+    input_vars = {i for i in graph_inputs(costs) if not isinstance(i, Constant)}
+    extra_vars_and_values = {
+        var: initial_point[var.name]
+        for var in model.value_vars
+        if var in input_vars and var not in grad_vars
+    }
+    return ValueGradFunction(
+        costs,
+        grad_vars,
+        extra_vars_and_values,
+        model=model,
+        initial_point=initial_point,
+        ravel_inputs=ravel_inputs,
+        **kwargs,
+    )
+
+
 class ValueGradFunction:
     """Create a PyTensor function that computes a value and its gradient.
 
@@ -479,53 +506,7 @@ def logp_dlogp_function(
             Compute the tempered logp `free_logp + alpha * observed_logp`.
             `alpha` can be changed using `ValueGradFunction.set_weights([alpha])`.
         """
-        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}")
-
-        if initial_point is None:
-            initial_point = self.initial_point(0)
-
-        # The compiled function does not depend on the initial_point values (those only seed
-        # the runtime-settable extra variables), so it is cached across calls with any point.
-        fn = self._logp_dlogp_function(
-            tuple(grad_vars),
-            tempered=tempered,
-            ravel_inputs=ravel_inputs,
-            initial_point=initial_point,
-            **kwargs,
-        )
-        fn.set_extra_values(initial_point)
-        return fn
-
-    def _logp_dlogp_function(
-        self, grad_vars, *, tempered=False, ravel_inputs=None, initial_point, **kwargs
-    ):
-        grad_vars = list(grad_vars)
-        if tempered:
-            costs = [self.varlogp, self.datalogp]
-        else:
-            costs = [self.logp()]
-
-        input_vars = {i for i in graph_inputs(costs) if not isinstance(i, Constant)}
-        extra_vars_and_values = {
-            var: initial_point[var.name]
-            for var in self.value_vars
-            if var in input_vars and var not in grad_vars
-        }
-        return ValueGradFunction(
-            costs,
-            grad_vars,
-            extra_vars_and_values,
-            model=self,
-            initial_point=initial_point,
-            ravel_inputs=ravel_inputs,
-            **kwargs,
-        )
+        raise NotImplementedError  # pragma: no cover
 
     def compile_logp(
         self,
@@ -938,12 +919,7 @@ def initial_point(self, random_seed: SeedSequenceSeed = None) -> dict[str, np.nd
         ip : dict of {str : array_like}
             Maps names of transformed variables to numeric initial values in the transformed space.
         """
-        fn = self._make_initial_point()
-        return Point(fn(random_seed), model=self)
-
-    def _make_initial_point(self):
-        # Compiled function takes the seed as an argument, so the cache is seed-independent.
-        return make_initial_point_fn(model=self, return_transformed=True)
+        raise NotImplementedError  # pragma: no cover
 
     def set_data(
         self,
@@ -1240,7 +1216,7 @@ def compile_fn(
         -------
         Compiled PyTensor function
         """
-        raise NotImplementedError
+        raise NotImplementedError  # pragma: no cover
 
     def profile(
         self,
@@ -1774,6 +1750,48 @@ class Model(BaseModel):
 
     """
 
+    def initial_point(self, random_seed: SeedSequenceSeed = None) -> dict[str, np.ndarray]:
+        fn = self._initial_point_fn()
+        return Point(fn(random_seed), model=self)
+
+    def _initial_point_fn(self, *, overrides=None, jitter_rvs=None, return_transformed=True):
+        return make_initial_point_fn(
+            model=self,
+            overrides=overrides,
+            jitter_rvs=jitter_rvs,
+            return_transformed=return_transformed,
+        )
+
+    def logp_dlogp_function(
+        self,
+        grad_vars=None,
+        tempered=False,
+        initial_point: PointType | None = None,
+        ravel_inputs: bool | None = None,
+        **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}")
+
+        if initial_point is None:
+            initial_point = self.initial_point(0)
+
+        fn = _make_value_grad_function(
+            self,
+            grad_vars,
+            tempered=tempered,
+            ravel_inputs=ravel_inputs,
+            initial_point=initial_point,
+            **kwargs,
+        )
+        fn.set_extra_values(initial_point)
+        return fn
+
     @overload
     def compile_fn(
         self,
@@ -2214,12 +2232,55 @@ def __init__(self, *args, **kwargs):
     logp = locally_cachedmethod(BaseModel.logp)
     dlogp = locally_cachedmethod(BaseModel.dlogp)
     d2logp = locally_cachedmethod(BaseModel.d2logp)
+
+    def initial_point(self, random_seed: SeedSequenceSeed = None) -> dict[str, np.ndarray]:
+        fn = self._initial_point_fn()
+        return Point(fn(random_seed), model=self)
+
+    @locally_cachedmethod
+    def _initial_point_fn(self, *, overrides=None, jitter_rvs=None, return_transformed=True):
+        # The compiled function takes the seed as an argument, so this does not depend on it.
+        return make_initial_point_fn(
+            model=self,
+            overrides=overrides,
+            jitter_rvs=jitter_rvs,
+            return_transformed=return_transformed,
+        )
+
+    def logp_dlogp_function(
+        self,
+        grad_vars=None,
+        tempered=False,
+        initial_point: PointType | None = None,
+        ravel_inputs: bool | None = None,
+        **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}")
+
+        if initial_point is None:
+            initial_point = self.initial_point(0)
+
+        fn = self._value_grad_function(
+            tuple(grad_vars),
+            tempered=tempered,
+            ravel_inputs=ravel_inputs,
+            initial_point=initial_point,
+            **kwargs,
+        )
+        fn.set_extra_values(initial_point)
+        return fn
+
     # The initial point only seeds the runtime-settable extra variables, so it is not part
-    # of the cache key (it is re-applied by `logp_dlogp_function` on every call).
-    _logp_dlogp_function = locally_cachedmethod(
-        BaseModel._logp_dlogp_function, ignore=("initial_point",)
-    )
-    _make_initial_point = locally_cachedmethod(BaseModel._make_initial_point)
+    # of the cache key: `logp_dlogp_function` re-applies it on every call.
+    @locally_cachedmethod(ignore=("initial_point",))
+    def _value_grad_function(self, grad_vars, **kwargs):
+        return _make_value_grad_function(self, grad_vars, **kwargs)
 
     @overload
     def compile_fn(

pymc/util.py

@@ -25,6 +25,7 @@
 
 from cachetools import LRUCache, cachedmethod
 from pytensor.compile import SharedVariable
+from pytensor.compile.io import In, Out
 from pytensor.graph.basic import Variable
 from xarray import Dataset, DataTree
 
@@ -302,6 +303,13 @@ def hashable(a=None) -> int:
         # lists are mutable and not hashable by default
         # for memoization, we need the hash to depend on the items
         return hash(tuple(hashable(i) for i in a))
+    if isinstance(a, set | frozenset):
+        # same as for lists, but order-insensitive
+        return hash(frozenset(hashable(i) for i in a))
+    if isinstance(a, In | Out):
+        # these wrap a variable with compilation options and are hashed by identity,
+        # so hash what they hold instead
+        return hashable(a.__dict__)
     try:
         return hash(a)
     except TypeError:
@@ -321,18 +329,28 @@ def hash_key(*args, **kwargs):
 
 
 class HashableWrapper:
-    __slots__ = ("obj",)
+    __slots__ = ("_hash", "obj")
 
     def __init__(self, obj):
         self.obj = obj
+        self._hash = hashable(obj)
 
     def __hash__(self):
         """Return a hash of the object."""
-        return hashable(self.obj)
+        return self._hash
 
     def __eq__(self, other):
-        """Compare this object with `other`."""
-        return self.obj == other
+        """Compare this object with `other`.
+
+        Compares the types and the hashes computed by :func:`hashable`, since the wrapped
+        objects may not support equality that returns a bool (arrays, or containers holding
+        them).
+        """
+        return (
+            isinstance(other, HashableWrapper)
+            and type(self.obj) is type(other.obj)
+            and self._hash == other._hash
+        )
 
     def __repr__(self):
         """Return a string representation of the object."""

tests/model/test_core.py

@@ -1380,7 +1380,7 @@ def test_logp_dlogp_function_is_cached(self):
         f1 = fm.logp_dlogp_function(ravel_inputs=True)
         f2 = fm.logp_dlogp_function(ravel_inputs=True)
         assert f1 is f2
-        assert "_logp_dlogp_function" in fm._cache
+        assert "_value_grad_function" in fm._cache
 
     def test_logp_dlogp_d2logp_graphs_are_cached(self):
         # Memoized graph construction returns the same object, so a freshly requested logp
@@ -1438,10 +1438,38 @@ def test_initial_point_is_cached(self):
 
         fm = freeze_model(m)
         ip1 = fm.initial_point(0)
-        assert "_make_initial_point" in fm._cache
+        assert "_initial_point_fn" in fm._cache
         np.testing.assert_allclose(ip1["x"], fm.initial_point(0)["x"])
         np.testing.assert_allclose(ip1["x"], m.initial_point(0)["x"])  # matches unfrozen
 
+    def test_repeated_sampling_does_not_recompile(self):
+        with pm.Model() as m:
+            x = pm.Normal("x", 0, 1, size=2)
+            pm.Normal("y", x, 1, observed=[0.3, -0.5])
+
+        sample_kwargs = {
+            "draws": 5,
+            "tune": 5,
+            "chains": 1,
+            "progressbar": False,
+            "nuts_sampler": "pymc",
+            "compute_convergence_checks": False,
+        }
+        fm = freeze_model(m)
+        with fm:
+            pm.sample(random_seed=0, **sample_kwargs)
+
+        n_compiles = [0]
+        orig_function = pytensor.function
+
+        def counting_function(*args, **kwargs):
+            n_compiles[0] += 1
+            return orig_function(*args, **kwargs)
+
+        with patch("pytensor.function", counting_function), fm:
+            pm.sample(random_seed=1, **sample_kwargs)
+        assert n_compiles[0] == 0
+
 
 def test_model_parent_set_programmatically():
     with pm.Model() as model: