pymc.model.transform.optimization.freeze_dims_and_data#

pymc.model.transform.optimization.freeze_dims_and_data(model, dims=None, data=None)[source]#

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:

modelModel: The model where to freeze dims and data.
dimsSequence of str, optional: The dimensions to freeze. If None, all dimensions are frozen. Pass an empty list to avoid freezing any dimension.
dataSequence 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.

Examples

import pymc as pm
import pytensor.tensor as pt

from pymc.model.transform.optimization 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)