pymc.model.transform.freeze_model

pymc.model.transform.freeze_model(model)

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 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

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)