# Copyright 2024 - present The PyMC Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections.abc import Sequence
from typing import cast
from pytensor.compile import SharedVariable
from pytensor.graph import Constant, FunctionGraph, Variable
from pytensor.graph.replace import clone_replace
from pytensor.graph.traversal import ancestors
from pymc.model.core import FrozenModel, Model
from pymc.model.fgraph import ModelFreeRV, fgraph_from_model, model_from_fgraph
def _constant_from_shared(shared: SharedVariable) -> Constant:
return shared.type.constant_type(type=shared.type, data=shared.get_value(), name=shared.name)
def _extract_initial_values(model: Model) -> dict[str, object]:
"""Return the model's non-default initial values, keyed by variable name.
Symbolic initial values reference variables of the model's graph, which the fgraph
round-trip clones, so they cannot be transplanted onto the rebuilt model and are
rejected.
"""
initial_values = {}
for rv, initval in model.rvs_to_initial_values.items():
if initval is None:
continue
if isinstance(initval, Variable) and not isinstance(initval, Constant):
raise NotImplementedError(
f"{rv.name} has a symbolic initial value, which cannot be transplanted onto "
"the transformed model. Only None, strategy strings and constant initial "
"values are supported."
)
initial_values[rv.name] = initval
return initial_values
[docs]
def freeze_dims_and_data(
model: Model, dims: Sequence[str] | None = None, data: Sequence[str] | None = None
) -> Model:
"""Recreate a Model with fixed RV dimensions and Data values.
The dimensions of the pre-existing RVs will no longer follow changes to the coordinates.
Likewise, it will not be possible to update pre-existing Data in the new model.
Note that any new RVs and Data created after calling this function will still be "unfrozen".
This transformation may allow more performant sampling, or compiling model functions to backends that
are more restrictive about dynamic shapes such as JAX.
Parameters
----------
model : Model
The model where to freeze dims and data.
dims : Sequence of str, optional
The dimensions to freeze.
If None, all dimensions are frozen. Pass an empty list to avoid freezing any dimension.
data : Sequence of str, optional
The data to freeze.
If None, all data are frozen. Pass an empty list to avoid freezing any data.
Returns
-------
Model
A new model with the specified dimensions and data frozen.
Notes
-----
Constant and strategy-string initial values are preserved on the new model. Symbolic
initial values (which reference variables of the original graph) are not supported.
Examples
--------
.. code-block:: python
import pymc as pm
import pytensor.tensor as pt
from pymc.model.transform import freeze_dims_and_data
with pm.Model() as m:
x = pm.Data("x", [0, 1, 2] * 1000)
y = pm.Normal("y", mu=pt.unique(x).mean())
# pt.unique(x).mean() has to be computed in every logp function evaluation
print("Logp eval time (1000x): ", m.profile(m.logp()).fct_call_time)
# pt.uniqe(x).mean() is cached in the logp function
frozen_m = freeze_dims_and_data(m)
print("Logp eval time (1000x): ", frozen_m.profile(frozen_m.logp()).fct_call_time)
"""
# fgraph_from_model does not carry initial values through the round-trip and rejects
# models that have them. Preserve them here: clear them for the round-trip and transplant
# them back onto the new model (matched by variable name) below.
initial_values = _extract_initial_values(model)
saved_initial_values = dict(model.rvs_to_initial_values)
try:
for rv in model.rvs_to_initial_values:
model.rvs_to_initial_values[rv] = None
fg, memo = fgraph_from_model(model)
finally:
model.rvs_to_initial_values.update(saved_initial_values)
if dims is None:
dims = tuple(model.dim_lengths.keys())
if data is None:
data = tuple(model.named_vars.keys())
# Replace mutable dim lengths and data by constants
frozen_replacements = {
memo[dim_length]: _constant_from_shared(dim_length)
for dim_length in (model.dim_lengths[dim_name] for dim_name in dims)
if isinstance(dim_length, SharedVariable)
}
frozen_replacements |= {
memo[datum].owner.inputs[0]: _constant_from_shared(datum)
for datum in (model.named_vars[datum_name] for datum_name in data)
if isinstance(datum, SharedVariable)
}
old_outs, old_coords, old_dim_lenghts = fg.outputs, fg._coords, fg._dim_lengths # type: ignore[attr-defined]
# Rebuild strict will force the recreation of RV nodes with updated static types
new_outs = clone_replace(old_outs, replace=frozen_replacements, rebuild_strict=False) # type: ignore[arg-type]
fg = FunctionGraph(outputs=new_outs, clone=False)
fg._coords = old_coords # type: ignore[attr-defined]
fg._dim_lengths = { # type: ignore[attr-defined]
dim: frozen_replacements.get(dim_length, dim_length)
for dim, dim_length in old_dim_lenghts.items()
}
# Recreate value variables from new RVs to propagate static types to logp graphs
replacements = {}
for node in fg.apply_nodes:
if not isinstance(node.op, ModelFreeRV):
continue
rv, old_value, *_ = node.inputs
transform = node.op.transform
if transform is None:
new_value = rv.type()
else:
new_value = transform.forward(rv, *rv.owner.inputs).type() # type: ignore[arg-type]
new_value.name = old_value.name
replacements[old_value] = new_value
fg.replace_all(tuple(replacements.items()), import_missing=True)
new_model = model_from_fgraph(fg, mutate_fgraph=True)
for name, initval in initial_values.items():
new_model.set_initval(new_model[name], initval)
return new_model
[docs]
def freeze_model(model: Model) -> FrozenModel:
"""Return a frozen copy of the model that caches its compiled functions.
On the frozen model, compiled functions (``compile_fn``, ``logp_dlogp_function``,
``initial_point``, and the forward-sampling function used by
``sample_prior_predictive`` / ``sample_posterior_predictive``) are compiled once and
reused across calls, so e.g. batched posterior predictive over changing ``pm.set_data``
values, or repeated ``pm.sample``, do not recompile. Seeding is re-applied on every
call, so cached functions stay reproducible.
To keep the cache valid the frozen model cannot be mutated: graph-mutating methods
(``register_rv``, ``add_coord``, ``set_initval``, ...) raise, and the dims and data
that any free variable depends on are frozen to constants as in
:func:`freeze_dims_and_data`. Data (and dims) that only Deterministics and observed
variables depend on remain updatable through ``pm.set_data`` — values and shapes are
runtime inputs of the cached functions, so updates and resizes take effect without
recompilation.
Functions with random variables compiled to backends that detach their RNGs at compile
time (JAX, MLX, PyTorch) cannot be reseeded and are compiled fresh on each call.
Constant and strategy-string initial values are preserved on the frozen model;
symbolic initial values are not supported.
Examples
--------
.. code-block:: python
import pymc as pm
from pymc.model.transform.optimization import freeze_model
with pm.Model() as m:
x = pm.Data("x", [0.0, 1.0, 2.0])
beta = pm.Normal("beta")
pm.Normal("y", mu=beta * x, observed=[1.0, 2.0, 3.0], shape=x.shape)
idata = pm.sample()
with freeze_model(m):
for x_batch in x_batches:
pm.set_data({"x": x_batch})
# Compiles on the first call only
pm.sample_posterior_predictive(idata, predictions=True)
"""
free_rv_ancestors = set(ancestors(model.free_RVs))
frozen_dims = [
name
for name, length in model.dim_lengths.items()
if isinstance(length, SharedVariable) and length in free_rv_ancestors
]
frozen_data = [
name
for name, var in model.named_vars.items()
if isinstance(var, SharedVariable) and var in free_rv_ancestors
]
frozen_model = freeze_dims_and_data(model, dims=frozen_dims, data=frozen_data)
# Retype the rebuilt model in place as a FrozenModel. This is the standard idiom for
# converting an instance to a sibling class: both are pure-Python subclasses of
# BaseModel with the same instance layout, so only the method resolution changes
# (mutators become unavailable, compiled functions become cached).
frozen_model.__class__ = FrozenModel # type: ignore[assignment]
return cast(FrozenModel, frozen_model)
__all__ = ("freeze_dims_and_data", "freeze_model")
