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..81289b1234 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._initial_point_fn( 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._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 77d16a8944..ba46e50101 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 # 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( 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/model/test_core.py b/tests/model/test_core.py index 7b99bf97d6..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,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: 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):