From 69d0df668ff02943901e5ac2813daf81bfd67790 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Thu, 30 Jul 2026 20:08:23 +0300 Subject: [PATCH 1/4] Fix cache keys that raise or never match `pymc.util.hashable` and `HashableWrapper` back the memoization helpers, and mishandle three kinds of values: - Looking up a key compared the wrapped objects, which raises `ValueError: truth value of an array is ambiguous` for arrays and for the containers holding them. Keys are now compared by type and by the hash `hashable` computes, so a cache can be keyed on a dict of arrays. - Sets fell through to being pickled whole, which is not stable across calls, so equal sets got different hashes and silently defeated the cache. They are now hashed by their elements, like lists and dicts. - `In`/`Out` wrap a variable with compilation options but are hashed by identity, and callers rebuild them on every call, so a key holding one never matched. They are now hashed by what they hold. Co-Authored-By: Claude Fable 5 --- pymc/util.py | 26 ++++++++++++++++++++++---- tests/test_util.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/pymc/util.py b/pymc/util.py index cbeaacfc0b..44a11e637f 100644 --- a/pymc/util.py +++ b/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.""" diff --git a/tests/test_util.py b/tests/test_util.py index c849515732..83c399c2c9 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -15,10 +15,12 @@ import arviz import numpy as np +import pytensor.tensor as pt import pytest import xarray from cachetools import cached +from pytensor.compile.io import In, Out import pymc as pm @@ -117,6 +119,39 @@ def test_hashing_of_rv_tuples(): assert isinstance(hashable(structure), int) +def test_hashable_of_sets(): + # Sets used to fall through to being pickled as a whole, which is not stable across + # calls, so equal sets got different hashes. + with pm.Model(): + rv = pm.Normal("rv") + + assert hashable({rv}) == hashable({rv}) + assert hashable({rv}) != hashable(set()) + assert hashable({1, 2}) == hashable({2, 1}) # order-insensitive + assert hashable(frozenset({1, 2})) == hashable(frozenset({1, 2})) + + +def test_hashable_of_function_inputs_and_outputs(): + # In/Out wrap a variable with compilation options and are hashed by identity, but + # callers rebuild them on every call. + x = pt.vector("x") + assert hashable(In(x, borrow=True)) == hashable(In(x, borrow=True)) + assert hashable(In(x, borrow=True)) != hashable(In(x, borrow=False)) + assert hashable(Out(x, borrow=True)) == hashable(Out(x, borrow=True)) + + +def test_hash_key_of_arrays(): + # Looking a key up compares it with the stored one, which must not raise for arrays or + # for the containers holding them. + key = hash_key({"x": np.zeros(3)}) + same = hash_key({"x": np.zeros(3)}) + other = hash_key({"x": np.ones(3)}) + + assert key == same + assert key != other + assert {key: "value"}[same] == "value" + + def test_hash_key(): class Bad1: def __hash__(self): From b8618a8a27c00bd4acc1ab71c7dd1b5c75175d63 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Thu, 30 Jul 2026 20:19:00 +0300 Subject: [PATCH 2/4] Reuse the initial point and trace functions when sampling a frozen model Repeated `pm.sample` on a frozen model still recompiled the initial-point and trace functions, so the caching a frozen model advertises never fired on the path that matters most. Both call sites now go through the model, which the cache-key fixes make possible. Repeated sampling of a frozen model now compiles nothing. Co-Authored-By: Claude Fable 5 --- pymc/backends/base.py | 7 ++++--- pymc/initial_point.py | 6 ++---- pymc/model/core.py | 9 +++++++-- tests/model/test_core.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 41 insertions(+), 9 deletions(-) diff --git a/pymc/backends/base.py b/pymc/backends/base.py index 39c6ea90ba..b18bd8547c 100644 --- a/pymc/backends/base.py +++ b/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 diff --git a/pymc/initial_point.py b/pymc/initial_point.py index 5530c232dd..0c21df10b6 100644 --- a/pymc/initial_point.py +++ b/pymc/initial_point.py @@ -104,8 +104,7 @@ 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._make_initial_point( overrides=overrides, jitter_rvs=jitter_rvs, return_transformed=True, @@ -113,8 +112,7 @@ def make_initial_point_fns_per_chain( ] * chains elif len(overrides) == chains: ipfns = [ - make_initial_point_fn( - model=model, + model._make_initial_point( jitter_rvs=jitter_rvs, overrides=chain_overrides, return_transformed=True, diff --git a/pymc/model/core.py b/pymc/model/core.py index 77d16a8944..56cd7cbe86 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -941,9 +941,14 @@ def initial_point(self, random_seed: SeedSequenceSeed = None) -> dict[str, np.nd fn = self._make_initial_point() return Point(fn(random_seed), model=self) - def _make_initial_point(self): + def _make_initial_point(self, *, overrides=None, jitter_rvs=None, return_transformed=True): # Compiled function takes the seed as an argument, so the cache is seed-independent. - return make_initial_point_fn(model=self, return_transformed=True) + return make_initial_point_fn( + model=self, + overrides=overrides, + jitter_rvs=jitter_rvs, + return_transformed=return_transformed, + ) def set_data( self, diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 7b99bf97d6..b51dc1c950 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -1442,6 +1442,34 @@ def test_initial_point_is_cached(self): 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: From ea49dafcfbb2828c0c4b5fe21969a7e6b644034a Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Sun, 2 Aug 2026 16:06:11 +0300 Subject: [PATCH 3/4] Keep the remaining caching in FrozenModel `logp_dlogp_function` and `initial_point` were still built around private methods in the base model that only existed to give the frozen model's cache something to wrap, the same shape that `compile_fn` was cleaned up from. Declare both in the base model and implement them in each final instead. `Model` builds its function every call, `FrozenModel` builds it through a cached method it owns. The part worth sharing, assembling the `ValueGradFunction`, moves to a module level helper both call, so the duplication is the argument handling rather than the logic. Co-Authored-By: Claude Fable 5 --- pymc/initial_point.py | 4 +- pymc/model/core.py | 182 +++++++++++++++++++++++++-------------- tests/model/test_core.py | 4 +- 3 files changed, 123 insertions(+), 67 deletions(-) diff --git a/pymc/initial_point.py b/pymc/initial_point.py index 0c21df10b6..81289b1234 100644 --- a/pymc/initial_point.py +++ b/pymc/initial_point.py @@ -104,7 +104,7 @@ def make_initial_point_fns_per_chain( # One strategy for all chains # Only one function compilation is needed. ipfns = [ - model._make_initial_point( + model._initial_point_fn( overrides=overrides, jitter_rvs=jitter_rvs, return_transformed=True, @@ -112,7 +112,7 @@ def make_initial_point_fns_per_chain( ] * chains elif len(overrides) == chains: ipfns = [ - model._make_initial_point( + model._initial_point_fn( jitter_rvs=jitter_rvs, overrides=chain_overrides, return_transformed=True, diff --git a/pymc/model/core.py b/pymc/model/core.py index 56cd7cbe86..8aa5ec715c 100644 --- a/pymc/model/core.py +++ b/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 def compile_logp( self, @@ -938,17 +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, *, overrides=None, jitter_rvs=None, return_transformed=True): - # Compiled function takes the seed as an argument, so the cache is seed-independent. - return make_initial_point_fn( - model=self, - overrides=overrides, - jitter_rvs=jitter_rvs, - return_transformed=return_transformed, - ) + raise NotImplementedError def set_data( self, @@ -1779,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, @@ -2219,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( diff --git a/tests/model/test_core.py b/tests/model/test_core.py index b51dc1c950..03ed48a556 100644 --- a/tests/model/test_core.py +++ b/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,7 +1438,7 @@ 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 From 8e7ec5acfcebe4a64845987e30ac85ad470f067f Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Sun, 2 Aug 2026 16:25:45 +0300 Subject: [PATCH 4/4] Exclude the abstract model methods from coverage Both finals override all three, and the base model cannot be instantiated, so the stubs are unreachable by construction and only show up as uncovered lines. Same marker the stand-in ops in `logprob` use. Co-Authored-By: Claude Fable 5 --- pymc/model/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 8aa5ec715c..ba46e50101 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -506,7 +506,7 @@ def logp_dlogp_function( Compute the tempered logp `free_logp + alpha * observed_logp`. `alpha` can be changed using `ValueGradFunction.set_weights([alpha])`. """ - raise NotImplementedError + raise NotImplementedError # pragma: no cover def compile_logp( self, @@ -919,7 +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. """ - raise NotImplementedError + raise NotImplementedError # pragma: no cover def set_data( self, @@ -1216,7 +1216,7 @@ def compile_fn( ------- Compiled PyTensor function """ - raise NotImplementedError + raise NotImplementedError # pragma: no cover def profile( self,