diff --git a/docs/source/api/model/core.rst b/docs/source/api/model/core.rst index afe588d46d..e15ef6b129 100644 --- a/docs/source/api/model/core.rst +++ b/docs/source/api/model/core.rst @@ -5,6 +5,8 @@ Model creation and inspection .. autosummary:: :toctree: generated/ + BaseModel + FrozenModel Model modelcontext diff --git a/docs/source/api/model/optimization.rst b/docs/source/api/model/optimization.rst index 8781c0953d..ffd23cf89f 100644 --- a/docs/source/api/model/optimization.rst +++ b/docs/source/api/model/optimization.rst @@ -5,3 +5,4 @@ Model Optimization :toctree: generated/ freeze_dims_and_data + freeze_model diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 05043ca6ff..ff346f7c7e 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -37,7 +37,7 @@ import pymc -from pymc.model import Model, modelcontext +from pymc.model import BaseModel, modelcontext from pymc.progress_bar import CustomProgress, default_progress_theme from pymc.pytensorf import PointFunc, extract_obs_data from pymc.util import get_default_varnames @@ -134,7 +134,7 @@ def dict_to_dataset_drop_incompatible_coords( return ds -def find_observations(model: Model) -> dict[str, Var]: +def find_observations(model: BaseModel) -> dict[str, Var]: """If there are observations available, return them as a dictionary.""" observations = {} for obs in model.observed_RVs: @@ -151,7 +151,7 @@ def find_observations(model: Model) -> dict[str, Var]: return observations -def find_constants(model: Model) -> dict[str, Var]: +def find_constants(model: BaseModel) -> dict[str, Var]: """If there are constants available, return them as a dictionary.""" model_vars = model.basic_RVs + model.deterministics + model.potentials value_vars = set(model.rvs_to_values.values()) @@ -174,7 +174,7 @@ def find_constants(model: Model) -> dict[str, Var]: def patch_nutpie_idata( idata: DataTree, - model: Model, + model: BaseModel, sampling_time: float, ) -> None: """Fix up the ``DataTree`` returned by ``nutpie.sample`` in place. @@ -216,7 +216,7 @@ def patch_nutpie_idata( idata.posterior.attrs[k] = v -def coords_and_dims_for_inferencedata(model: Model) -> tuple[dict[str, Any], dict[str, Any]]: +def coords_and_dims_for_inferencedata(model: BaseModel) -> tuple[dict[str, Any], dict[str, Any]]: """Parse PyMC model coords and dims format to one accepted by InferenceData.""" coords = { cname: np.array(cvals) if isinstance(cvals, tuple) else cvals @@ -283,7 +283,7 @@ def insert(self, k: str, v, idx: int): class DataTreeConverter: """Encapsulate conventions of InferenceData schema in DataTree conversion.""" - model: Model | None = None + model: BaseModel | None = None posterior_predictive: Mapping[str, np.ndarray] | None = None predictions: Mapping[str, np.ndarray] | None = None prior: Mapping[str, np.ndarray] | None = None @@ -620,7 +620,7 @@ def to_inference_data( coords: CoordSpec | None = None, dims: DimSpec | None = None, sample_dims: list | None = None, - model: Model = None, + model: BaseModel = None, save_warmup: bool | None = None, include_transformed: bool = False, ) -> DataTree: @@ -652,7 +652,7 @@ def to_inference_data( Map of coordinate names to coordinate values dims : dict of {str: list of str}, optional Map of variable names to the coordinate names to use to index its dimensions. - model : Model, optional + model : BaseModel, optional Model used to generate ``trace``. It is not necessary to pass ``model`` if in ``with`` context. save_warmup : bool, optional @@ -689,7 +689,7 @@ def to_inference_data( def predictions_to_inference_data( predictions, posterior_trace: MultiTrace | None = None, - model: Model | None = None, + model: BaseModel | None = None, coords: CoordSpec | None = None, dims: DimSpec | None = None, sample_dims: list | None = None, @@ -709,7 +709,7 @@ def predictions_to_inference_data( ``pymc.sample_posterior_predictive``. Specifically, any variable whose shape is a deterministic function of the shape of any predictor (explanatory, independent, etc.) variables must be *removed* from this trace. - model: Model + model: BaseModel The pymc model. It can be omitted if within a model context. coords: Dict[str, array-like[Any]] Coordinates for the variables. Map from coordinate names to coordinate values. diff --git a/pymc/backends/zarr.py b/pymc/backends/zarr.py index 9f120de72f..1bc20b8475 100644 --- a/pymc/backends/zarr.py +++ b/pymc/backends/zarr.py @@ -31,7 +31,7 @@ ) from pymc.backends.base import BaseTrace from pymc.blocking import StatDtype, StatShape -from pymc.model.core import Model, modelcontext +from pymc.model.core import BaseModel, modelcontext from pymc.step_methods.compound import ( BlockedStep, CompoundStep, @@ -80,7 +80,7 @@ class ZarrChain(_ZarrChainBase, BaseTrace): said stats with the accompanying stepper index. synchronizer : zarr.sync.Synchronizer | None The synchronizer to use for the underlying zarr arrays. - model : Model + model : BaseModel If None, the model is taken from the `with` context. vars : Sequence[TensorVariable] | None Sampling values will be stored for these variables. If None, @@ -99,7 +99,7 @@ def __init__( store: BaseStore | MutableMapping, stats_bijection: StatsBijection, synchronizer: Synchronizer | None = None, - model: Model | None = None, + model: BaseModel | None = None, vars: Sequence[TensorVariable] | None = None, test_point: dict[str, np.ndarray] | None = None, draws_per_chunk: int = 1, @@ -412,7 +412,7 @@ def init_trace( draws: int, tune: int, step: BlockedStep | CompoundStep, - model: Model | None = None, + model: BaseModel | None = None, vars: Sequence[TensorVariable] | None = None, test_point: dict[str, np.ndarray] | None = None, ): diff --git a/pymc/data.py b/pymc/data.py index 41d88d0d35..7f184ef8f5 100644 --- a/pymc/data.py +++ b/pymc/data.py @@ -38,7 +38,7 @@ from pymc.vartypes import isgenerator if typing.TYPE_CHECKING: - from pymc.model.core import Model + from pymc.model.core import BaseModel __all__ = [ "Data", @@ -226,7 +226,7 @@ def Data( dims: Sequence[str] | None = None, coords: dict[str, Sequence | np.ndarray] | None = None, infer_dims_and_coords=False, - model: Union["Model", None] = None, + model: Union["BaseModel", None] = None, **kwargs, ) -> SharedVariable | TensorConstant: """Create a data container that registers a data variable with the model. @@ -291,7 +291,7 @@ def Data( ... model.set_data("data", data_vals) ... idatas.append(pm.sample()) """ - from pymc.model.core import modelcontext + from pymc.model.core import Model, modelcontext if coords is None: coords = {} @@ -307,6 +307,8 @@ def Data( "No model on context stack, which is needed to instantiate a data container. " "Add variable inside a 'with model:' block." ) + if not isinstance(model, Model): + raise TypeError(f"Cannot add new data to an immutable {type(model).__name__}.") name = model.name_for(name) # Transform `value` it to something digestible for PyTensor. diff --git a/pymc/dims/model.py b/pymc/dims/model.py index 6ef5a4373d..7f7ae2e4c5 100644 --- a/pymc/dims/model.py +++ b/pymc/dims/model.py @@ -23,13 +23,13 @@ convert_dims, convert_dims_with_ellipsis, ) +from pymc.model.core import BaseModel, Model, modelcontext from pymc.model.core import Deterministic as RegularDeterministic -from pymc.model.core import Model, modelcontext from pymc.model.core import Potential as RegularPotential def Data( - name: str, value, dims: Dims = None, model: Model | None = None, **kwargs + name: str, value, dims: Dims = None, model: BaseModel | None = None, **kwargs ) -> XTensorSharedVariable: """Wrapper around pymc.Data that returns an XtensorVariable. @@ -37,6 +37,8 @@ def Data( These are always forwarded to the model object. """ model = modelcontext(model) + if not isinstance(model, Model): + raise TypeError(f"Cannot add new data to an immutable {type(model).__name__}.") dims = convert_dims(dims) # type: ignore[assignment] value = xtensor_shared(value, dims=dims, **kwargs, name=name) model.register_data_var(value, dims=value.dims) diff --git a/pymc/model/core.py b/pymc/model/core.py index 676c5d6f67..c796a6048c 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -57,11 +57,13 @@ from pymc.pytensorf import ( PointFunc, SeedSequenceSeed, + collect_default_updates, compile, convert_observed_data, gradient, hessian, join_nonshared_inputs, + reseed_rngs, resolve_backend_compile_kwargs, rewrite_pregrad, ) @@ -72,6 +74,8 @@ get_transformed_name, get_value_vars_from_user_vars, get_var_name, + locally_cachedmethod, + makeiter, treedict, treelist, ) @@ -80,7 +84,9 @@ sparse = lazy_scipy_module("sparse") __all__ = [ + "BaseModel", "Deterministic", + "FrozenModel", "Model", "Point", "Potential", @@ -101,15 +107,15 @@ class ModelManager(threading.local): """ def __init__(self): - self.active_contexts: list[Model] = [] + self.active_contexts: list[BaseModel] = [] @property - def current_context(self) -> Model | None: + def current_context(self) -> BaseModel | None: """Return the innermost context of any current contexts.""" return self.active_contexts[-1] if self.active_contexts else None @property - def parent_context(self) -> Model | None: + def parent_context(self) -> BaseModel | None: """Return the parent context to the active context, if any.""" return self.active_contexts[-2] if len(self.active_contexts) > 1 else None @@ -119,10 +125,10 @@ def parent_context(self) -> Model | None: MODEL_MANAGER = ModelManager() -def modelcontext(model: Model | None) -> Model: +def modelcontext(model: BaseModel | None) -> BaseModel: """Return the given model or, if None was supplied, try to find one in the context stack.""" if model is None: - model = Model.get_context(error_if_none=False) + model = BaseModel.get_context(error_if_none=False) if model is None: # TODO: This should be a ValueError, but that breaks @@ -309,125 +315,30 @@ def __call__(cls: type[Model], *args, **kwargs): return instance -class Model(WithMemoization, metaclass=ContextMeta): - """Encapsulates the variables and likelihood factors of a model. - - Model class can be used for creating class based models. To create - a class based model you should inherit from :class:`~pymc.Model` and - override the `__init__` method with arbitrary definitions (do not - forget to call base class :meth:`pymc.Model.__init__` first). - - Parameters - ---------- - name : str - name that will be used as prefix for names of all random - variables defined within model - coords : dict - Xarray-like coordinate keys and values. These coordinates can be used - to specify the shape of random variables and to label (but not specify) - the shape of Determinsitic, Potential and Data objects. - Other than specifying the shape of random variables, coordinates have no - effect on the model. They can't be used for label-based broadcasting or indexing. - You must use numpy-like operations for those behaviors. - check_bounds : bool - Ensure that input parameters to distributions are in a valid - range. If your model is built in a way where you know your - parameters can only take on valid values you can set this to - False for increased speed. This should not be used if your model - contains discrete variables. - model : PyMC model, optional - A parent model that this model belongs to. If not specified and the current model - is created inside another model's context, the parent model will be set to that model. - If `None` the model will not have a parent. - - Examples - -------- - Use context manager to define model and respective variables - - .. code-block:: python - - import pymc as pm - - with pm.Model() as model: - x = pm.Normal("x") - - - Use object API to define model and respective variables - - .. code-block:: python - - import pymc as pm - - model = pm.Model() - x = pm.Normal("x", model=model) - - - Use coords for defining the shape of random variables and labeling other model variables - - .. code-block:: python - - import pymc as pm - import numpy as np - - coords = { - "feature": ["A", "B", "C"], - "trial": [1, 2, 3, 4, 5], - } - - with pm.Model(coords=coords) as model: - # Variable will have default dim label `intercept__dim_0` - intercept = pm.Normal("intercept", shape=(3,)) - # Variable will have shape (3,) and dim label `feature` - beta = pm.Normal("beta", dims=("feature",)) - - # Dims below are only used for labeling, they have no effect on shape - # Variable will have default dim label `idx__dim_0` - idx = pm.Data("idx", np.array([0, 1, 1, 2, 2])) - x = pm.Data("x", np.random.normal(size=(5, 3)), dims=("trial", "feature")) - # single dim can be passed as string - mu = pm.Deterministic("mu", intercept[idx] + beta @ x, dims="trial") - - # Dims controls the shape of the variable - # If not specified, it would be inferred from the shape of the observations - y = pm.Normal("y", mu=mu, observed=[-1, 0, 0, 1, 1], dims=("trial",)) - - - Define nested models, and provide name for variable name prefixing - - .. code-block:: python - - import pymc as pm - - with pm.Model(name="root") as root: - x = pm.Normal("x") # Variable wil be named "root::x" - - with pm.Model(name="first") as first: - # Variable will belong to root and first - y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y" - - # Can pass parent model explicitly - with pm.Model(name="second", model=root) as second: - # Variable will belong to root and second - z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" - - # Set None for standalone model - with pm.Model(name="third", model=None) as third: - # Variable will belong to third only - w = pm.Normal("w") # Variable wil be named "third::w" +def _rng_detaching_linker(mode) -> bool: + """Whether ``mode``'s linker copies RNG shared variables at compile time. + Such compiled functions cannot be reseeded afterwards, so functions with RNGs must be + seeded before compilation and are not served from the compilation cache. + """ + from pytensor.compile.mode import get_mode + from pytensor.link.jax.linker import JAXLinker + from pytensor.link.mlx.linker import MLXLinker + from pytensor.link.pytorch.linker import PytorchLinker - Set `check_bounds` to False for models with only continuous variables and default transformers - PyMC will remove the bounds check from the model logp which can speed up sampling - - .. code-block:: python - - import pymc as pm + try: + linker = get_mode(mode).linker + except Exception: + return False + return isinstance(linker, JAXLinker | MLXLinker | PytorchLinker) - with pm.Model(check_bounds=False) as model: - sigma = pm.HalfNormal("sigma") - x = pm.Normal("x", sigma=sigma) # No bounds check will be performed on `sigma` +class BaseModel(WithMemoization, metaclass=ContextMeta): + """Functionality shared by mutable and frozen models. + Abstract: instantiate :class:`~pymc.Model` (mutable), or create a + :class:`~pymc.FrozenModel` with + :func:`pymc.model.transform.optimization.freeze_model`. """ def __enter__(self): @@ -475,6 +386,8 @@ def __init__( *, model: _UnsetType | None | Model = UNSET, ): + if type(self) is BaseModel: + raise TypeError("BaseModel is abstract; instantiate Model instead.") self.name = self._validate_name(name) self.check_bounds = check_bounds self._parent = model if not isinstance(model, _UnsetType) else MODEL_MANAGER.parent_context @@ -519,7 +432,7 @@ def __init__( @classmethod def get_context( cls, error_if_none: bool = True, allow_block_model_access: bool = False - ) -> Model | None: + ) -> BaseModel | None: model = MODEL_MANAGER.current_context if isinstance(model, BlockModelAccess) and not allow_block_model_access: raise BlockModelAccessError(model.error_msg_on_access) @@ -569,14 +482,27 @@ def logp_dlogp_function( 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 only depends on the graph and compile args, not on the + # initial_point values (those just seed the runtime-settable extra variables), so + # cache the compilation and re-apply this call's extra values. + fn = self._logp_dlogp_function( + tuple(grad_vars), tempered=tempered, ravel_inputs=ravel_inputs, **kwargs + ) + fn.set_extra_values(initial_point) + return fn + + def _logp_dlogp_function(self, grad_vars, *, tempered=False, ravel_inputs=None, **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)} - if initial_point is None: - initial_point = self.initial_point(0) + initial_point = self.initial_point(0) extra_vars_and_values = { var: initial_point[var.name] for var in self.value_vars @@ -681,6 +607,10 @@ def logp( ) -> Variable | list[Variable]: """Elemwise log-probability of the model. + On frozen models the returned graph is memoized (keyed on the arguments), so + repeated calls return the same object and ``compile_fn`` can reuse a compiled + function. Do not mutate it in place. + Parameters ---------- vars : list of random variables or potential terms, optional @@ -951,74 +881,6 @@ def shape_from_dims(self, dims): shape.extend(np.shape(self.coords[dim])) return tuple(shape) - def add_coord( - self, - name: str, - values: Sequence | np.ndarray | None = None, - *, - length: int | Variable | None = None, - ): - """Register a dimension coordinate with the model. - - Parameters - ---------- - name : str - Name of the dimension. - Forbidden: {"chain", "draw", "__sample__"} - values : optional, array_like - Coordinate values or ``None`` (for auto-numbering). - If ``None`` is passed, a ``length`` must be specified. - length : optional, scalar - A scalar of the dimensions length. - Defaults to ``pytensor.tensor.constant(len(values))``. - """ - if name in {"draw", "chain", "__sample__"}: - raise ValueError( - "Dimensions can not be named `draw`, `chain` or `__sample__`, " - "as those are reserved for use in `InferenceData`." - ) - if self.named_vars.tree_contains(name): - raise ValueError( - f"Dimension name '{name}' conflicts with an existing model variable name." - ) - if values is None and length is None: - raise ValueError( - f"Either `values` or `length` must be specified for the '{name}' dimension." - ) - if values is not None: - # Conversion to a tuple ensures that the coordinate values are immutable. - # Also unlike numpy arrays the's tuple.index(...) which is handy to work with. - values = tuple(values) - if name in self.coords: - if not np.array_equal(values, self.coords[name]): - raise ValueError(f"Duplicate and incompatible coordinate: {name}.") - return - if length is not None and not isinstance(length, int | Variable): - raise ValueError( - f"The `length` passed for the '{name}' coord must be an int, PyTensor Variable or None." - ) - if length is None: - length = len(values) - if not isinstance(length, Variable): - length = pytensor.shared(length, name=name) - assert length.type.ndim == 0 - self._dim_lengths[name] = length - self._coords[name] = values - - def add_coords( - self, - coords: dict[str, Sequence | None], - *, - lengths: dict[str, int | Variable | None] | None = None, - ): - """Vectorized version of ``Model.add_coord``.""" - if coords is None: - return - lengths = lengths or {} - - for name, values in coords.items(): - self.add_coord(name, values, length=lengths.get(name, None)) - def set_dim(self, name: str, new_length: int, coord_values: Sequence | None = None): """Update a mutable dimension. @@ -1067,16 +929,12 @@ 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 = make_initial_point_fn(model=self, return_transformed=True) + fn = self._make_initial_point() return Point(fn(random_seed), model=self) - def set_initval(self, rv_var, initval): - """Set an initial value (strategy) for a random variable.""" - if initval is not None and not isinstance(initval, Variable | str): - # Convert scalars or array-like inputs to ndarrays - initval = rv_var.type.filter(initval) - - self.rvs_to_initial_values[rv_var] = initval + 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) def set_data( self, @@ -1102,10 +960,15 @@ def set_data( """ shared_object = self[name] if not isinstance(shared_object, SharedVariable): + frozen_hint = ( + " The model is frozen and only data that no free variable depends on can be set." + if isinstance(self, FrozenModel) + else "" + ) raise TypeError( f"The variable `{name}` must be a `SharedVariable` " "(created through `pm.Data()` to allow updating.) " - f"The current type is: {type(shared_object)}" + f"The current type is: {type(shared_object)}.{frozen_hint}" ) if isinstance(values, list): @@ -1213,476 +1076,208 @@ def set_data( shared_object.set_value(values) - def register_rv( - self, - rv_var: RandomVariable, - name: str, - *, - observed=None, - total_size=None, - dims=None, - default_transform=UNSET, - transform=UNSET, - initval=None, - ) -> TensorVariable: - """Register an (un)observed random variable with the model. - - Parameters - ---------- - rv_var : TensorVariable - name : str - Intended name for the model variable. - observed : array_like, optional - Data values for observed variables. - total_size : scalar - upscales logp of variable with ``coef = total_size/var.shape[0]`` - dims : tuple - Dimension names for the variable. - default_transform - A default transform for the random variable in log-likelihood space. - transform - Additional transform which may be applied after default transform. - initval - The initial value of the random variable. + @property + def prefix(self) -> str: + if self.isroot or not self.parent.prefix: + name = self.name + else: + name = f"{self.parent.prefix}::{self.name}" + return name - Returns - ------- - TensorVariable - """ - name = self.name_for(name) - rv_var.name = name + def name_for(self, name): + """Check if name has prefix and adds if needed.""" + name = self._validate_name(name) + if self.prefix: + if not name.startswith(self.prefix + "::"): + return f"{self.prefix}::{name}" + else: + return name + else: + return name - # Associate previously unknown dimension names with - # the length of the corresponding RV dimension. - if dims is not None: - for d, dname in enumerate(dims): - if not isinstance(dname, str): - raise TypeError(f"Dims must be string. Got {dname} of type {type(dname)}") - if dname not in self.dim_lengths: - self.add_coord(dname, values=None, length=rv_var.shape[d]) - - if observed is None: - if total_size is not None: - raise ValueError("total_size can only be passed to observed RVs") - self.free_RVs.append(rv_var) - self.create_value_var(rv_var, transform=transform, default_transform=default_transform) - self.add_named_variable(rv_var, dims) - self.set_initval(rv_var, initval) + def name_of(self, name): + """Check if name has prefix and deletes if needed.""" + name = self._validate_name(name) + if not self.prefix or not name: + return name + elif name.startswith(self.prefix + "::"): + return name[len(self.prefix) + 2 :] else: - if total_size is None and isinstance(observed, TensorVariable): - for node in ancestors([observed]): - if node.owner is not None and isinstance(node.owner.op, MinibatchOp): - warnings.warn( - f"total_size not provided for observed variable `{name}` that uses pm.Minibatch" - ) - break - if not is_valid_observed(observed): - raise TypeError( - "Variables that depend on other nodes cannot be used for observed data." - f"The data variable was: {observed}" - ) + return name - # `rv_var` is potentially changed by `make_obs_var`, - # for example into a new graph for imputation of missing data. - rv_var = self.make_obs_var( - rv_var, observed, dims, default_transform, transform, total_size - ) + def __getitem__(self, key): + """Get the variable named `key`.""" + try: + return self.named_vars[key] + except KeyError as e: + try: + return self.named_vars[self.name_for(key)] + except KeyError: + raise e - return rv_var + def __contains__(self, key): + """Check if the model contains a variable named `key`.""" + return key in self.named_vars or self.name_for(key) in self.named_vars - def make_obs_var( - self, - rv_var: TensorVariable, - data: np.ndarray, - dims, - default_transform: Transform | None, - transform: Transform | None, - total_size: int | None, - ) -> TensorVariable: - """Create a `TensorVariable` for an observed random variable. + def __copy__(self): + """Clone the model.""" + return self.copy() - Parameters - ---------- - rv_var : TensorVariable - The random variable that is observed. - Its dimensionality must be compatible with the data already. - data : array_like - The observed data. - dims : tuple - Dimension names for the variable. - default_transform - A transform for the random variable in log-likelihood space. - transform - Additional transform which may be applied after default transform. + def __deepcopy__(self, _): + """Clone the model.""" + return self.copy() - Returns - ------- - TensorVariable + def copy(self): """ - name = rv_var.name - data = convert_observed_data(data).astype(rv_var.dtype) - - if data.ndim != rv_var.ndim: - raise ShapeError( - "Dimensionality of data and RV don't match.", actual=data.ndim, expected=rv_var.ndim - ) - - mask = getattr(data, "mask", None) - if mask is not None: - impute_message = ( - f"Data in {rv_var} contains missing values and" - " will be automatically imputed from the" - " sampling distribution." - ) - warnings.warn(impute_message, ImputationWarning) + Clone the model. - if total_size is not None: - raise ValueError("total_size is not compatible with imputed variables") + To access variables in the cloned model use `cloned_model["var_name"]`. - from pymc.distributions.distribution import create_partial_observed_rv + Examples + -------- + .. code-block:: python - ( - (observed_rv, observed_mask), - (unobserved_rv, _), - joined_rv, - ) = create_partial_observed_rv(rv_var, mask) - observed_data = pt.as_tensor(data.data[observed_mask]) + import pymc as pm + import copy - # Register ObservedRV corresponding to observed component - observed_rv.name = f"{name}_observed" - self.create_value_var( - observed_rv, transform=transform, default_transform=None, value_var=observed_data - ) - self.add_named_variable(observed_rv) - self.observed_RVs.append(observed_rv) + with pm.Model() as m: + p = pm.Beta("p", 1, 1) + x = pm.Bernoulli("x", p=p, shape=(3,)) - # Register FreeRV corresponding to unobserved components - self.register_rv( - unobserved_rv, - f"{name}_unobserved", - transform=transform, - default_transform=default_transform, - ) + clone_m = copy.copy(m) - # Register Deterministic that combines observed and missing - # Note: This can widely increase memory consumption during sampling for large datasets - rv_var = Deterministic(name, joined_rv, self, dims) + # Access cloned variables by name + clone_x = clone_m["x"] - else: - if sparse.issparse(data): - import pytensor.sparse as pt_sparse + # z will be part of clone_m but not m + z = pm.Deterministic("z", clone_x + 1) + """ + from pymc.model.fgraph import clone_model - data = pt_sparse.basic.as_sparse(data, name=name) - elif not isinstance(data, Variable): - data = pt.as_tensor_variable(data, name=name) + return clone_model(self) - if total_size: - from pymc.variational.minibatch_rv import create_minibatch_rv + def replace_rvs_by_values( + self, + graphs: Sequence[TensorVariable], + **kwargs, + ) -> list[TensorVariable]: + """Clone and replace random variables in graphs with their value variables. - rv_var = create_minibatch_rv(rv_var, total_size) - rv_var.name = name + This will *not* recompute test values in the resulting graphs. - rv_var.tag.observations = data - self.create_value_var( - rv_var, transform=transform, default_transform=None, value_var=data - ) - self.add_named_variable(rv_var, dims) - self.observed_RVs.append(rv_var) + Parameters + ---------- + graphs : array_like + The graphs in which to perform the replacements. - return rv_var + Returns + ------- + array_like + """ + return replace_rvs_by_values( + graphs, + rvs_to_values=self.rvs_to_values, + rvs_to_transforms=self.rvs_to_transforms, + ) - def create_value_var( + @overload + def compile_fn( self, - rv_var: TensorVariable, + outs: Variable | Sequence[Variable], *, - default_transform: Transform, - transform: Transform, - value_var: Variable | None = None, - ) -> TensorVariable: - """Create a ``TensorVariable`` that will be used as the random variable's "value" in log-likelihood graphs. + inputs: Sequence[Variable] | None = None, + mode=None, + point_fn: Literal[True] = True, + **kwargs, + ) -> PointFunc: ... - In general, we'll call this type of variable the "value" variable. + @overload + def compile_fn( + self, + outs: Variable | Sequence[Variable], + *, + inputs: Sequence[Variable] | None = None, + mode=None, + point_fn: Literal[False], + **kwargs, + ) -> Function: ... - In all other cases, the role of the value variable is taken by - observed data. That's why value variables are only referenced in - this branch of the conditional. + def compile_fn( + self, + outs: Variable | Sequence[Variable], + *, + inputs: Sequence[Variable] | None = None, + mode=None, + point_fn: bool = True, + **kwargs, + ) -> PointFunc | Function: + """Compiles a PyTensor function. Parameters ---------- - rv_var : TensorVariable - - default_transform: Transform - A transform for the random variable in log-likelihood space. - - transform: Transform - Additional transform which may be applied after default transform. - - value_var : Variable, optional + outs : Variable or sequence of Variables + PyTensor variable or iterable of PyTensor variables. + inputs : sequence of Variables, optional + PyTensor input variables, Required if there is more than one input. + mode + PyTensor compilation mode, default=None. + point_fn : bool + Whether to wrap the compiled function in a PointFunc, which takes a Point + dictionary with model variable names and values as input. + Other keyword arguments : + Any other keyword argument is sent to :py:func:`pymc.pytensorf.compile`. Returns ------- - TensorVariable + Compiled PyTensor function """ - from pymc.distributions.transforms import ChainedTransform, _default_transform + if inputs is None: + inputs = list(explicit_graph_inputs(outs)) + if (not point_fn) and len(inputs) > 1: + raise ValueError( + "compile_fn requires inputs to be specified when there is more than one input and point_fn is disabled." + ) - # Make the value variable a transformed value variable, - # if there's an applicable transform - if transform is None and default_transform is UNSET: - default_transform = None - warnings.warn( - "To disable default transform, please use default_transform=None" - " instead of transform=None. Setting transform to None will" - " not have any effect in future.", - UserWarning, - ) - - if default_transform is UNSET: - if rv_var.owner is None: - default_transform = None - else: - default_transform = _default_transform(rv_var.owner.op, rv_var) - - if transform is UNSET: - transform = default_transform - elif transform is not None and default_transform is not None: - transform = ChainedTransform([default_transform, transform]) - - if value_var is None: - if transform is None: - # Create value variable with the same type as the RV - value_var = rv_var.type() - value_var.name = rv_var.name - else: - # Create value variable with the same type as the transformed RV - value_var = transform.forward(rv_var, *rv_var.owner.inputs).type() - value_var.name = f"{rv_var.name}_{transform.name}__" - value_var.tag.transform = transform - - self.rvs_to_transforms[rv_var] = transform - self.rvs_to_values[rv_var] = value_var - self.values_to_rvs[value_var] = rv_var - - return value_var - - def register_data_var(self, data, dims=None): - """Register a data variable with the model.""" - self.data_vars.append(data) - self.add_named_variable(data, dims=dims) - - def add_named_variable(self, var, dims: tuple[str | None, ...] | None = None): - """Add a random graph variable to the named variables of the model. - - This can include several types of variables such basic_RVs, Data, Deterministics, - and Potentials. - - Parameters - ---------- - var - - dims : tuple, optional - - """ - if var.name is None: - raise ValueError("Variable is unnamed.") - if self.named_vars.tree_contains(var.name): - raise ValueError(f"Variable name {var.name} already exists.") - if var.name in self.coords: - raise ValueError( - f"Variable name '{var.name}' conflicts with an existing dimension name." - ) - - if dims is not None: - if isinstance(dims, str): - dims = (dims,) - for dim in dims: - if dim not in self.coords and dim is not None: - raise ValueError(f"Dimension {dim} is not specified in `coords`.") - # This check implicitly states that only vars with .ndim attribute can have dims - if var.ndim != len(dims): - raise ValueError( - f"{var} has {var.ndim} dims but {len(dims)} dim labels were provided." - ) - self.named_vars_to_dims[var.name] = dims - - self.named_vars[var.name] = var - if not hasattr(self, self.name_of(var.name)): - setattr(self, self.name_of(var.name), var) - - @property - def prefix(self) -> str: - if self.isroot or not self.parent.prefix: - name = self.name - else: - name = f"{self.parent.prefix}::{self.name}" - return name - - def name_for(self, name): - """Check if name has prefix and adds if needed.""" - name = self._validate_name(name) - if self.prefix: - if not name.startswith(self.prefix + "::"): - return f"{self.prefix}::{name}" - else: - return name - else: - return name - - def name_of(self, name): - """Check if name has prefix and deletes if needed.""" - name = self._validate_name(name) - if not self.prefix or not name: - return name - elif name.startswith(self.prefix + "::"): - return name[len(self.prefix) + 2 :] + random_seed = kwargs.pop("random_seed", None) + + if _rng_detaching_linker(mode) and collect_default_updates( + inputs=list(inputs), outputs=list(makeiter(outs)) + ): + # These linkers detach RNGs at compile time, so a compiled function cannot be + # reseeded. Seed before compiling and skip the cache. + kwargs.setdefault("allow_input_downcast", True) + kwargs.setdefault("accept_inplace", True) + with self: + fn = compile(inputs, outs, mode=mode, random_seed=random_seed, **kwargs) else: - return name - - def __getitem__(self, key): - """Get the variable named `key`.""" - try: - return self.named_vars[key] - except KeyError as e: - try: - return self.named_vars[self.name_for(key)] - except KeyError: - raise e - - def __contains__(self, key): - """Check if the model contains a variable named `key`.""" - return key in self.named_vars or self.name_for(key) in self.named_vars - - def __copy__(self): - """Clone the model.""" - return self.copy() - - def __deepcopy__(self, _): - """Clone the model.""" - return self.copy() - - def copy(self): - """ - Clone the model. - - To access variables in the cloned model use `cloned_model["var_name"]`. - - Examples - -------- - .. code-block:: python - - import pymc as pm - import copy - - with pm.Model() as m: - p = pm.Beta("p", 1, 1) - x = pm.Bernoulli("x", p=p, shape=(3,)) - - clone_m = copy.copy(m) - - # Access cloned variables by name - clone_x = clone_m["x"] - - # z will be part of clone_m but not m - z = pm.Deterministic("z", clone_x + 1) - """ - from pymc.model.fgraph import clone_model - - return clone_model(self) - - def replace_rvs_by_values( - self, - graphs: Sequence[TensorVariable], - **kwargs, - ) -> list[TensorVariable]: - """Clone and replace random variables in graphs with their value variables. - - This will *not* recompute test values in the resulting graphs. - - Parameters - ---------- - graphs : array_like - The graphs in which to perform the replacements. - - Returns - ------- - array_like - """ - return replace_rvs_by_values( - graphs, - rvs_to_values=self.rvs_to_values, - rvs_to_transforms=self.rvs_to_transforms, - ) - - @overload - def compile_fn( - self, - outs: Variable | Sequence[Variable], - *, - inputs: Sequence[Variable] | None = None, - mode=None, - point_fn: Literal[True] = True, - **kwargs, - ) -> PointFunc: ... + # Compile seed-independently (cached on frozen models) and reseed on each call, + # in the order pytensorf.compile uses, so a cached function yields the same RNG + # stream as a fresh compile. + fn, rng_updates = self._compile_fn(outs, inputs=tuple(inputs), mode=mode, **kwargs) + if rng_updates: + reseed_rngs(rng_updates, random_seed) - @overload - def compile_fn( - self, - outs: Variable | Sequence[Variable], - *, - inputs: Sequence[Variable] | None = None, - mode=None, - point_fn: Literal[False], - **kwargs, - ) -> Function: ... + if point_fn: + return PointFunc(fn) + return fn - def compile_fn( + def _compile_fn( self, outs: Variable | Sequence[Variable], *, inputs: Sequence[Variable] | None = None, mode=None, - point_fn: bool = True, **kwargs, - ) -> PointFunc | Function: - """Compiles a PyTensor function. - - Parameters - ---------- - outs : Variable or sequence of Variables - PyTensor variable or iterable of PyTensor variables. - inputs : sequence of Variables, optional - PyTensor input variables, Required if there is more than one input. - mode - PyTensor compilation mode, default=None. - point_fn : bool - Whether to wrap the compiled function in a PointFunc, which takes a Point - dictionary with model variable names and values as input. - Other keyword arguments : - Any other keyword argument is sent to :py:func:`pymc.pytensorf.compile`. - - Returns - ------- - Compiled PyTensor function - """ - if inputs is None: - inputs = list(explicit_graph_inputs(outs)) - if (not point_fn) and len(inputs) > 1: - raise ValueError( - "compile_fn requires inputs to be specified when there is more than one input and point_fn is disabled." - ) - + ) -> tuple[Function, list[SharedVariable]]: + # random_seed=False keeps the compiled function seed-independent and cacheable; + # compile_fn reseeds it afterwards. The RNG variables are collected once and cached + # with the function so they need not be re-walked on each call. + kwargs.setdefault("allow_input_downcast", True) + kwargs.setdefault("accept_inplace", True) with self: - fn = compile( - inputs, - outs, - allow_input_downcast=True, - accept_inplace=True, - mode=mode, - **kwargs, - ) - - if point_fn: - return PointFunc(fn) - return fn + fn = compile(inputs, outs, mode=mode, random_seed=False, **kwargs) + rng_updates = collect_default_updates(inputs=list(inputs), outputs=list(makeiter(outs))) + return fn, list(rng_updates) def profile( self, @@ -2085,6 +1680,536 @@ def table( ) +class Model(BaseModel): + """Encapsulates the variables and likelihood factors of a model. + + Model class can be used for creating class based models. To create + a class based model you should inherit from :class:`~pymc.Model` and + override the `__init__` method with arbitrary definitions (do not + forget to call base class :meth:`pymc.Model.__init__` first). + + Parameters + ---------- + name : str + name that will be used as prefix for names of all random + variables defined within model + coords : dict + Xarray-like coordinate keys and values. These coordinates can be used + to specify the shape of random variables and to label (but not specify) + the shape of Determinsitic, Potential and Data objects. + Other than specifying the shape of random variables, coordinates have no + effect on the model. They can't be used for label-based broadcasting or indexing. + You must use numpy-like operations for those behaviors. + check_bounds : bool + Ensure that input parameters to distributions are in a valid + range. If your model is built in a way where you know your + parameters can only take on valid values you can set this to + False for increased speed. This should not be used if your model + contains discrete variables. + model : PyMC model, optional + A parent model that this model belongs to. If not specified and the current model + is created inside another model's context, the parent model will be set to that model. + If `None` the model will not have a parent. + + Examples + -------- + Use context manager to define model and respective variables + + .. code-block:: python + + import pymc as pm + + with pm.Model() as model: + x = pm.Normal("x") + + + Use object API to define model and respective variables + + .. code-block:: python + + import pymc as pm + + model = pm.Model() + x = pm.Normal("x", model=model) + + + Use coords for defining the shape of random variables and labeling other model variables + + .. code-block:: python + + import pymc as pm + import numpy as np + + coords = { + "feature": ["A", "B", "C"], + "trial": [1, 2, 3, 4, 5], + } + + with pm.Model(coords=coords) as model: + # Variable will have default dim label `intercept__dim_0` + intercept = pm.Normal("intercept", shape=(3,)) + # Variable will have shape (3,) and dim label `feature` + beta = pm.Normal("beta", dims=("feature",)) + + # Dims below are only used for labeling, they have no effect on shape + # Variable will have default dim label `idx__dim_0` + idx = pm.Data("idx", np.array([0, 1, 1, 2, 2])) + x = pm.Data("x", np.random.normal(size=(5, 3)), dims=("trial", "feature")) + # single dim can be passed as string + mu = pm.Deterministic("mu", intercept[idx] + beta @ x, dims="trial") + + # Dims controls the shape of the variable + # If not specified, it would be inferred from the shape of the observations + y = pm.Normal("y", mu=mu, observed=[-1, 0, 0, 1, 1], dims=("trial",)) + + + Define nested models, and provide name for variable name prefixing + + .. code-block:: python + + import pymc as pm + + with pm.Model(name="root") as root: + x = pm.Normal("x") # Variable wil be named "root::x" + + with pm.Model(name="first") as first: + # Variable will belong to root and first + y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y" + + # Can pass parent model explicitly + with pm.Model(name="second", model=root) as second: + # Variable will belong to root and second + z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" + + # Set None for standalone model + with pm.Model(name="third", model=None) as third: + # Variable will belong to third only + w = pm.Normal("w") # Variable wil be named "third::w" + + + Set `check_bounds` to False for models with only continuous variables and default transformers + PyMC will remove the bounds check from the model logp which can speed up sampling + + .. code-block:: python + + import pymc as pm + + with pm.Model(check_bounds=False) as model: + sigma = pm.HalfNormal("sigma") + x = pm.Normal("x", sigma=sigma) # No bounds check will be performed on `sigma` + + + Notes + ----- + A mutable model recompiles its functions on every call. To reuse compiled functions + across calls — e.g. batched posterior predictive over changing ``pm.set_data`` values, + or repeated ``pm.sample`` — freeze the model with + :func:`pymc.model.transform.optimization.freeze_model`. The frozen copy caches the + graphs and compiled functions it builds (``logp``, ``compile_fn``, + ``logp_dlogp_function``, ``initial_point``, and the forward-sampling function) and + forbids the mutations that would make them stale. + + """ + + def add_coord( + self, + name: str, + values: Sequence | np.ndarray | None = None, + *, + length: int | Variable | None = None, + ): + """Register a dimension coordinate with the model. + + Parameters + ---------- + name : str + Name of the dimension. + Forbidden: {"chain", "draw", "__sample__"} + values : optional, array_like + Coordinate values or ``None`` (for auto-numbering). + If ``None`` is passed, a ``length`` must be specified. + length : optional, scalar + A scalar of the dimensions length. + Defaults to ``pytensor.tensor.constant(len(values))``. + """ + if name in {"draw", "chain", "__sample__"}: + raise ValueError( + "Dimensions can not be named `draw`, `chain` or `__sample__`, " + "as those are reserved for use in `InferenceData`." + ) + if self.named_vars.tree_contains(name): + raise ValueError( + f"Dimension name '{name}' conflicts with an existing model variable name." + ) + if values is None and length is None: + raise ValueError( + f"Either `values` or `length` must be specified for the '{name}' dimension." + ) + if values is not None: + # Conversion to a tuple ensures that the coordinate values are immutable. + # Also unlike numpy arrays the's tuple.index(...) which is handy to work with. + values = tuple(values) + if name in self.coords: + if not np.array_equal(values, self.coords[name]): + raise ValueError(f"Duplicate and incompatible coordinate: {name}.") + return + if length is not None and not isinstance(length, int | Variable): + raise ValueError( + f"The `length` passed for the '{name}' coord must be an int, PyTensor Variable or None." + ) + if length is None: + length = len(values) + if not isinstance(length, Variable): + length = pytensor.shared(length, name=name) + assert length.type.ndim == 0 + self._dim_lengths[name] = length + self._coords[name] = values + + def add_coords( + self, + coords: dict[str, Sequence | None], + *, + lengths: dict[str, int | Variable | None] | None = None, + ): + """Vectorized version of ``Model.add_coord``.""" + if coords is None: + return + lengths = lengths or {} + + for name, values in coords.items(): + self.add_coord(name, values, length=lengths.get(name, None)) + + def set_initval(self, rv_var, initval): + """Set an initial value (strategy) for a random variable.""" + if initval is not None and not isinstance(initval, Variable | str): + # Convert scalars or array-like inputs to ndarrays + initval = rv_var.type.filter(initval) + + self.rvs_to_initial_values[rv_var] = initval + + def register_rv( + self, + rv_var: RandomVariable, + name: str, + *, + observed=None, + total_size=None, + dims=None, + default_transform=UNSET, + transform=UNSET, + initval=None, + ) -> TensorVariable: + """Register an (un)observed random variable with the model. + + Parameters + ---------- + rv_var : TensorVariable + name : str + Intended name for the model variable. + observed : array_like, optional + Data values for observed variables. + total_size : scalar + upscales logp of variable with ``coef = total_size/var.shape[0]`` + dims : tuple + Dimension names for the variable. + default_transform + A default transform for the random variable in log-likelihood space. + transform + Additional transform which may be applied after default transform. + initval + The initial value of the random variable. + + Returns + ------- + TensorVariable + """ + name = self.name_for(name) + rv_var.name = name + + # Associate previously unknown dimension names with + # the length of the corresponding RV dimension. + if dims is not None: + for d, dname in enumerate(dims): + if not isinstance(dname, str): + raise TypeError(f"Dims must be string. Got {dname} of type {type(dname)}") + if dname not in self.dim_lengths: + self.add_coord(dname, values=None, length=rv_var.shape[d]) + + if observed is None: + if total_size is not None: + raise ValueError("total_size can only be passed to observed RVs") + self.free_RVs.append(rv_var) + self.create_value_var(rv_var, transform=transform, default_transform=default_transform) + self.add_named_variable(rv_var, dims) + self.set_initval(rv_var, initval) + else: + if total_size is None and isinstance(observed, TensorVariable): + for node in ancestors([observed]): + if node.owner is not None and isinstance(node.owner.op, MinibatchOp): + warnings.warn( + f"total_size not provided for observed variable `{name}` that uses pm.Minibatch" + ) + break + if not is_valid_observed(observed): + raise TypeError( + "Variables that depend on other nodes cannot be used for observed data." + f"The data variable was: {observed}" + ) + + # `rv_var` is potentially changed by `make_obs_var`, + # for example into a new graph for imputation of missing data. + rv_var = self.make_obs_var( + rv_var, observed, dims, default_transform, transform, total_size + ) + + return rv_var + + def make_obs_var( + self, + rv_var: TensorVariable, + data: np.ndarray, + dims, + default_transform: Transform | None, + transform: Transform | None, + total_size: int | None, + ) -> TensorVariable: + """Create a `TensorVariable` for an observed random variable. + + Parameters + ---------- + rv_var : TensorVariable + The random variable that is observed. + Its dimensionality must be compatible with the data already. + data : array_like + The observed data. + dims : tuple + Dimension names for the variable. + default_transform + A transform for the random variable in log-likelihood space. + transform + Additional transform which may be applied after default transform. + + Returns + ------- + TensorVariable + """ + name = rv_var.name + data = convert_observed_data(data).astype(rv_var.dtype) + + if data.ndim != rv_var.ndim: + raise ShapeError( + "Dimensionality of data and RV don't match.", actual=data.ndim, expected=rv_var.ndim + ) + + mask = getattr(data, "mask", None) + if mask is not None: + impute_message = ( + f"Data in {rv_var} contains missing values and" + " will be automatically imputed from the" + " sampling distribution." + ) + warnings.warn(impute_message, ImputationWarning) + + if total_size is not None: + raise ValueError("total_size is not compatible with imputed variables") + + from pymc.distributions.distribution import create_partial_observed_rv + + ( + (observed_rv, observed_mask), + (unobserved_rv, _), + joined_rv, + ) = create_partial_observed_rv(rv_var, mask) + observed_data = pt.as_tensor(data.data[observed_mask]) + + # Register ObservedRV corresponding to observed component + observed_rv.name = f"{name}_observed" + self.create_value_var( + observed_rv, transform=transform, default_transform=None, value_var=observed_data + ) + self.add_named_variable(observed_rv) + self.observed_RVs.append(observed_rv) + + # Register FreeRV corresponding to unobserved components + self.register_rv( + unobserved_rv, + f"{name}_unobserved", + transform=transform, + default_transform=default_transform, + ) + + # Register Deterministic that combines observed and missing + # Note: This can widely increase memory consumption during sampling for large datasets + rv_var = Deterministic(name, joined_rv, self, dims) + + else: + if sparse.issparse(data): + import pytensor.sparse as pt_sparse + + data = pt_sparse.basic.as_sparse(data, name=name) + elif not isinstance(data, Variable): + data = pt.as_tensor_variable(data, name=name) + + if total_size: + from pymc.variational.minibatch_rv import create_minibatch_rv + + rv_var = create_minibatch_rv(rv_var, total_size) + rv_var.name = name + + rv_var.tag.observations = data + self.create_value_var( + rv_var, transform=transform, default_transform=None, value_var=data + ) + self.add_named_variable(rv_var, dims) + self.observed_RVs.append(rv_var) + + return rv_var + + def create_value_var( + self, + rv_var: TensorVariable, + *, + default_transform: Transform, + transform: Transform, + value_var: Variable | None = None, + ) -> TensorVariable: + """Create a ``TensorVariable`` that will be used as the random variable's "value" in log-likelihood graphs. + + In general, we'll call this type of variable the "value" variable. + + In all other cases, the role of the value variable is taken by + observed data. That's why value variables are only referenced in + this branch of the conditional. + + Parameters + ---------- + rv_var : TensorVariable + + default_transform: Transform + A transform for the random variable in log-likelihood space. + + transform: Transform + Additional transform which may be applied after default transform. + + value_var : Variable, optional + + Returns + ------- + TensorVariable + """ + from pymc.distributions.transforms import ChainedTransform, _default_transform + + # Make the value variable a transformed value variable, + # if there's an applicable transform + if transform is None and default_transform is UNSET: + default_transform = None + warnings.warn( + "To disable default transform, please use default_transform=None" + " instead of transform=None. Setting transform to None will" + " not have any effect in future.", + UserWarning, + ) + + if default_transform is UNSET: + if rv_var.owner is None: + default_transform = None + else: + default_transform = _default_transform(rv_var.owner.op, rv_var) + + if transform is UNSET: + transform = default_transform + elif transform is not None and default_transform is not None: + transform = ChainedTransform([default_transform, transform]) + + if value_var is None: + if transform is None: + # Create value variable with the same type as the RV + value_var = rv_var.type() + value_var.name = rv_var.name + else: + # Create value variable with the same type as the transformed RV + value_var = transform.forward(rv_var, *rv_var.owner.inputs).type() + value_var.name = f"{rv_var.name}_{transform.name}__" + value_var.tag.transform = transform + + self.rvs_to_transforms[rv_var] = transform + self.rvs_to_values[rv_var] = value_var + self.values_to_rvs[value_var] = rv_var + + return value_var + + def register_data_var(self, data, dims=None): + """Register a data variable with the model.""" + self.data_vars.append(data) + self.add_named_variable(data, dims=dims) + + def add_named_variable(self, var, dims: tuple[str | None, ...] | None = None): + """Add a random graph variable to the named variables of the model. + + This can include several types of variables such basic_RVs, Data, Deterministics, + and Potentials. + + Parameters + ---------- + var + + dims : tuple, optional + + """ + if var.name is None: + raise ValueError("Variable is unnamed.") + if self.named_vars.tree_contains(var.name): + raise ValueError(f"Variable name {var.name} already exists.") + if var.name in self.coords: + raise ValueError( + f"Variable name '{var.name}' conflicts with an existing dimension name." + ) + + if dims is not None: + if isinstance(dims, str): + dims = (dims,) + for dim in dims: + if dim not in self.coords and dim is not None: + raise ValueError(f"Dimension {dim} is not specified in `coords`.") + # This check implicitly states that only vars with .ndim attribute can have dims + if var.ndim != len(dims): + raise ValueError( + f"{var} has {var.ndim} dims but {len(dims)} dim labels were provided." + ) + self.named_vars_to_dims[var.name] = dims + + self.named_vars[var.name] = var + if not hasattr(self, self.name_of(var.name)): + setattr(self, self.name_of(var.name), var) + + +class FrozenModel(BaseModel): + """A model whose graph is immutable and whose compiled functions are cached. + + Create one with :func:`pymc.model.transform.optimization.freeze_model`; this class + cannot be instantiated directly. Graph-mutating methods do not exist on it, which is + what makes caching safe without any invalidation. ``set_data`` remains available for + data that no free variable depends on — values and shapes are runtime inputs of the + cached functions. + """ + + def __init__(self, *args, **kwargs): + raise TypeError( + "FrozenModel cannot be instantiated directly. " + "Create one from an existing model with freeze_model." + ) + + # Cached graphs and compiled functions. The seed is applied outside `_compile_fn` + # (in `compile_fn`) and taken as an argument by the initial-point function, so all + # cached entries are seed-independent. + logp = locally_cachedmethod(BaseModel.logp) + dlogp = locally_cachedmethod(BaseModel.dlogp) + d2logp = locally_cachedmethod(BaseModel.d2logp) + _logp_dlogp_function = locally_cachedmethod(BaseModel._logp_dlogp_function) + _make_initial_point = locally_cachedmethod(BaseModel._make_initial_point) + _compile_fn = locally_cachedmethod(BaseModel._compile_fn) + + class BlockModelAccess(Model): """Can be used to prevent user access to Model contexts.""" diff --git a/pymc/model/transform/__init__.py b/pymc/model/transform/__init__.py index b56729481c..1dea85a61b 100644 --- a/pymc/model/transform/__init__.py +++ b/pymc/model/transform/__init__.py @@ -28,13 +28,14 @@ extract_deterministics, insert_deterministics, ) -from pymc.model.transform.optimization import freeze_dims_and_data +from pymc.model.transform.optimization import freeze_dims_and_data, freeze_model __all__ = ( "change_value_transforms", "do", "extract_deterministics", "freeze_dims_and_data", + "freeze_model", "insert_deterministics", "observe", "prune_vars_detached_from_observed", diff --git a/pymc/model/transform/conditioning.py b/pymc/model/transform/conditioning.py index a7e2d011e5..046549e6c8 100644 --- a/pymc/model/transform/conditioning.py +++ b/pymc/model/transform/conditioning.py @@ -22,7 +22,7 @@ from pytensor.tensor import TensorVariable from pymc.logprob.transforms import Transform -from pymc.model.core import Model +from pymc.model.core import BaseModel, Model from pymc.model.fgraph import ( ModelDeterministic, ModelFreeRV, @@ -315,7 +315,7 @@ def change_value_transforms( def remove_value_transforms( - model: Model, + model: BaseModel, vars: Sequence[ModelVariable] | None = None, ) -> Model: r"""Remove the value variables transforms in the model. diff --git a/pymc/model/transform/optimization.py b/pymc/model/transform/optimization.py index 8537eed69c..1f91281649 100644 --- a/pymc/model/transform/optimization.py +++ b/pymc/model/transform/optimization.py @@ -12,12 +12,14 @@ # 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 +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 Model +from pymc.model.core import FrozenModel, Model from pymc.model.fgraph import ModelFreeRV, fgraph_from_model, model_from_fgraph @@ -25,6 +27,27 @@ 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 + + def freeze_dims_and_data( model: Model, dims: Sequence[str] | None = None, data: Sequence[str] | None = None ) -> Model: @@ -54,6 +77,10 @@ def freeze_dims_and_data( 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 -------- @@ -76,7 +103,17 @@ def freeze_dims_and_data( print("Logp eval time (1000x): ", frozen_m.profile(frozen_m.logp()).fct_call_time) """ - fg, memo = fgraph_from_model(model) + # 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()) @@ -120,7 +157,74 @@ def freeze_dims_and_data( replacements[old_value] = new_value fg.replace_all(tuple(replacements.items()), import_missing=True) - return model_from_fgraph(fg, mutate_fgraph=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 + + +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. -__all__ = ("freeze_dims_and_data",) + 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") diff --git a/pymc/model/transform_values.py b/pymc/model/transform_values.py index 91503cd5b7..035fe115df 100644 --- a/pymc/model/transform_values.py +++ b/pymc/model/transform_values.py @@ -21,19 +21,19 @@ from pytensor.graph.basic import Variable from pytensor.xtensor.vectorization import vectorize_graph as xvectorize_graph -from pymc.model import Model, modelcontext +from pymc.model import BaseModel, modelcontext from pymc.pytensorf import compile, replace_vars_in_graphs def _build_transform_graph( - model: Model, + model: BaseModel, forward: bool, ) -> tuple[list[Variable], list[Variable]]: """Build a per-sample graph that applies transforms to all free RVs. Parameters ---------- - model : Model + model : BaseModel PyMC model whose free RVs define the transforms. forward : bool ``True`` for natural → unconstrained, ``False`` for the inverse. @@ -69,7 +69,7 @@ def _eval_transform_graph( *, forward: bool, dataset: xr.Dataset, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), compile_kwargs: dict | None = None, ): @@ -133,7 +133,7 @@ def _eval_transform_graph( def unconstrain_values( values: xr.Dataset, *, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), compile_kwargs: dict | None = None, ) -> xr.Dataset: @@ -167,7 +167,7 @@ def unconstrain_values( def constrain_values( values: xr.Dataset, *, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), compile_kwargs: dict | None = None, ) -> xr.Dataset: diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py index 5347f38086..e91885f4b0 100644 --- a/pymc/pytensorf.py +++ b/pymc/pytensorf.py @@ -899,7 +899,7 @@ def resolve_backend_compile_kwargs(backend: str | None, compile_kwargs: dict | N def compile( inputs, outputs, - random_seed: SeedSequenceSeed = None, + random_seed: SeedSequenceSeed | bool = None, mode=None, **kwargs, ) -> Function: @@ -946,8 +946,9 @@ def compile( ], ) - # We always reseed random variables as this provides RNGs with no chances of collision - if rng_updates: + # Reseed for collision-free RNGs, unless random_seed=False decouples seeding from + # compilation so the function can be cached and reseeded by the caller. + if rng_updates and random_seed is not False: rngs = cast(list[SharedVariable], list(rng_updates)) reseed_rngs(rngs, random_seed) diff --git a/pymc/sampling/deterministic.py b/pymc/sampling/deterministic.py index 49c0b1d2b0..2351c5cb06 100644 --- a/pymc/sampling/deterministic.py +++ b/pymc/sampling/deterministic.py @@ -16,7 +16,7 @@ from xarray import Dataset, DataTree, merge from pymc.backends.arviz import apply_function_over_dataset, coords_and_dims_for_inferencedata -from pymc.model.core import Model, modelcontext +from pymc.model.core import BaseModel, modelcontext from pymc.pytensorf import resolve_backend_compile_kwargs @@ -24,7 +24,7 @@ def compute_deterministics( dataset: Dataset, *, var_names: Sequence[str] | None = None, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), merge_dataset: bool = False, progressbar: bool = True, @@ -40,7 +40,7 @@ def compute_deterministics( var_names : sequence of str, optional List of names of deterministic variable to compute. If None, compute all deterministics in the model. - model : Model, optional + model : BaseModel, optional Model to use. If None, use context model. sample_dims : sequence of str, default ("chain", "draw") Sample (batch) dimensions of the dataset over which to compute the deterministics. diff --git a/pymc/sampling/forward.py b/pymc/sampling/forward.py index dc7f6f0af3..4b52064dc4 100644 --- a/pymc/sampling/forward.py +++ b/pymc/sampling/forward.py @@ -46,7 +46,7 @@ from pymc.blocking import PointType from pymc.distributions.shape_utils import change_dist_size from pymc.exceptions import ImplicitFreezeWarning -from pymc.model import Model, modelcontext +from pymc.model import BaseModel, modelcontext from pymc.progress_bar import create_simple_progress, default_progress_theme from pymc.pytensorf import compile, resolve_backend_compile_kwargs, rvs_in_graph from pymc.util import ( @@ -73,7 +73,7 @@ def _build_constant_data( trace_constant_data: dict[str, np.ndarray], trace_coords: dict[str, np.ndarray], - model: Model, + model: BaseModel, ) -> dict[Variable, Any]: """Collect trace-time values for data-like nodes in the model. @@ -194,7 +194,7 @@ def _maybe_warn_implicit_freeze( rv_set: set[Variable], vars_in_trace_set: set[Variable], constant_data: dict[Variable, Any], - model: Model, + model: BaseModel, trace_coords: dict[str, np.ndarray], ) -> None: """Warn when an auto-frozen trace RV has a volatile ancestor the user may have wanted resampled.""" @@ -266,6 +266,7 @@ def compile_forward_sampling_function( constant_data: dict[Variable, Any] | None = None, volatile_vars: set[Variable] | None = None, freeze_vars: set[Variable] | None = None, + model: BaseModel | None = None, **kwargs, ) -> tuple[Callable[..., np.ndarray | list[np.ndarray]], set[Variable]]: """Compile a function to draw samples, conditioned on the values of some variables. @@ -378,8 +379,17 @@ def expand(var): # the entire graph list(walk(fg.outputs, expand)) + if model is None: + fn = compile(inputs, fg.outputs, on_unused_input="ignore", **kwargs) + else: + # Route through the model so the compiled function is cached and reused across calls + # (e.g. batched posterior predictive with changing pm.set_data). + fn = model.compile_fn( + fg.outputs, inputs=inputs, point_fn=False, on_unused_input="ignore", **kwargs + ) + return ( - compile(inputs, fg.outputs, on_unused_input="ignore", **kwargs), + fn, set(basic_rvs) & volatile_closure, # Basic RVs that will be resampled ) @@ -456,7 +466,7 @@ def draw( return [np.stack(v) for v in drawn_values] -def observed_dependent_deterministics(model: Model, extra_observeds=None): +def observed_dependent_deterministics(model: BaseModel, extra_observeds=None): """Find deterministics that depend directly on observed variables.""" if extra_observeds is None: extra_observeds = [] @@ -474,7 +484,7 @@ def observed_dependent_deterministics(model: Model, extra_observeds=None): @overload def sample_prior_predictive( draws: int = 500, - model: Model | None = None, + model: BaseModel | None = None, var_names: Iterable[str] | None = None, random_seed: RandomState = None, return_inferencedata: Literal[True] = True, @@ -485,7 +495,7 @@ def sample_prior_predictive( @overload def sample_prior_predictive( draws: int = 500, - model: Model | None = None, + model: BaseModel | None = None, var_names: Iterable[str] | None = None, random_seed: RandomState = None, return_inferencedata: Literal[False] = False, @@ -495,7 +505,7 @@ def sample_prior_predictive( ) -> dict[str, np.ndarray]: ... def sample_prior_predictive( draws: int = 500, - model: Model | None = None, + model: BaseModel | None = None, var_names: Iterable[str] | None = None, random_seed: RandomState = None, return_inferencedata: bool = True, @@ -509,7 +519,7 @@ def sample_prior_predictive( ---------- draws : int Number of samples from the prior predictive to generate. Defaults to 500. - model : Model (optional if in ``with`` context) + model : BaseModel (optional if in ``with`` context) var_names : Iterable[str] A list of names of variables for which to compute the prior predictive samples. Defaults to both observed and unobserved RVs. Transformed values @@ -568,6 +578,7 @@ def sample_prior_predictive( vars_in_trace=[], basic_rvs=model.basic_RVs, random_seed=random_seed, + model=model, **compile_kwargs, ) @@ -595,7 +606,7 @@ def sample_prior_predictive( @overload def sample_posterior_predictive( trace, - model: Model | None = None, + model: BaseModel | None = None, *, var_names: str | list[str] | None = None, sample_vars: str | list[str] | None = None, @@ -614,7 +625,7 @@ def sample_posterior_predictive( @overload def sample_posterior_predictive( trace, - model: Model | None = None, + model: BaseModel | None = None, *, var_names: str | list[str] | None = None, sample_vars: str | list[str] | None = None, @@ -632,7 +643,7 @@ def sample_posterior_predictive( ) -> dict[str, np.ndarray]: ... def sample_posterior_predictive( trace, - model: Model | None = None, + model: BaseModel | None = None, *, var_names: str | list[str] | None = None, sample_vars: str | list[str] | None = None, @@ -663,7 +674,7 @@ def sample_posterior_predictive( trace : backend, list, Dataset, DataTree, or MultiTrace Trace generated from MCMC sampling, or a list of dicts (eg. points or from :func:`~pymc.find_MAP`), or :class:`xarray.Dataset` (eg. DataTree.posterior or DataTree.prior) - model : Model (optional if in ``with`` context) + model : BaseModel (optional if in ``with`` context) Model to be used to generate the posterior predictive samples. It will generally be the model used to generate the `trace`, but it doesn't need to be. sample_vars : str or list of str, optional @@ -1228,6 +1239,7 @@ def sample_posterior_predictive( constant_data=constant_data, volatile_vars=volatile_vars, freeze_vars=frozen_vars, + model=model, **compile_kwargs, ) sampler_fn = point_wrapper(_sampler_fn) diff --git a/pymc/stats/log_density.py b/pymc/stats/log_density.py index 554d8bcd0a..3de76e4664 100644 --- a/pymc/stats/log_density.py +++ b/pymc/stats/log_density.py @@ -20,7 +20,7 @@ apply_function_over_dataset, coords_and_dims_for_inferencedata, ) -from pymc.model import Model, modelcontext +from pymc.model import BaseModel, modelcontext from pymc.pytensorf import resolve_backend_compile_kwargs __all__ = ("compute_log_likelihood", "compute_log_prior") @@ -33,7 +33,7 @@ def compute_log_likelihood( *, var_names: Sequence[str] | None = None, extend_inferencedata: bool = True, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), progressbar=True, backend: str | None = None, @@ -50,7 +50,7 @@ def compute_log_likelihood( Defaults to all observed variables. extend_inferencedata : bool, default True Whether to extend the original InferenceData or return a new one - model : Model, optional + model : BaseModel, optional sample_dims : sequence of str, default ("chain", "draw") progressbar : bool, default True backend : str, optional @@ -82,7 +82,7 @@ def compute_log_prior( *, var_names: Sequence[str] | None = None, extend_inferencedata: bool = True, - model: Model | None = None, + model: BaseModel | None = None, sample_dims: Sequence[str] = ("chain", "draw"), progressbar=True, backend: str | None = None, @@ -99,7 +99,7 @@ def compute_log_prior( Defaults to all all free variables. extend_inferencedata : bool, default True Whether to extend the original InferenceData or return a new one - model : Model, optional + model : BaseModel, optional sample_dims : sequence of str, default ("chain", "draw") progressbar : bool, default True backend : str, optional @@ -131,7 +131,7 @@ def compute_log_density( *, var_names: Sequence[str] | None = None, extend_inferencedata: bool = True, - model: Model | None = None, + model: BaseModel | None = None, kind: Literal["likelihood", "prior"] = "likelihood", sample_dims: Sequence[str] = ("chain", "draw"), progressbar=True, @@ -150,7 +150,7 @@ def compute_log_density( Defaults to all all free variables. extend_inferencedata : bool, default True Whether to extend the original InferenceData or return a new one - model : Model, optional + model : BaseModel, optional kind: Literal["likelihood", "prior"] Whether to compute the log density of the observed random variables (likelihood) or to compute the log density of the latent random variables (prior). This diff --git a/pymc/util.py b/pymc/util.py index e144057461..a801c59451 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -356,6 +356,7 @@ def __setstate__(self, state): def locally_cachedmethod(f): + """Cache a method's return value on ``self._cache``, keyed by the call arguments.""" from collections import defaultdict def self_cache_fn(f_name): diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 7bcb3997bc..9200c51641 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -33,7 +33,8 @@ import scipy.stats as st from pytensor.compile.mode import get_default_mode -from pytensor.graph import graph_inputs +from pytensor.compile.sharedvalue import SharedVariable +from pytensor.graph import Constant, graph_inputs from pytensor.graph.traversal import get_var_by_name from pytensor.link.numba import NumbaLinker from pytensor.raise_op import Assert @@ -58,6 +59,8 @@ from pymc.logprob.basic import transformed_conditional_logp from pymc.logprob.transforms import IntervalTransform from pymc.model import Point, ValueGradFunction, modelcontext +from pymc.model.core import BaseModel, FrozenModel +from pymc.model.transform.optimization import freeze_model from pymc.pytensorf import floatX, inputvars from pymc.variational.minibatch_rv import MinibatchRandomVariable from tests.models import simple_model @@ -1177,6 +1180,232 @@ def test_compile_fn(): np.testing.assert_allclose(result_compute, result_expect) +class TestFrozenModelCaching: + def test_unfrozen_model_does_not_cache(self): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + mu = pm.Deterministic("mu", x * 2) + pm.Normal("y", x, 1, observed=[0.3, -0.5]) + + assert m.compile_fn(mu, inputs=[x], point_fn=False) is not m.compile_fn( + mu, inputs=[x], point_fn=False + ) + assert m.logp() is not m.logp() + assert m.logp_dlogp_function(ravel_inputs=True) is not m.logp_dlogp_function( + ravel_inputs=True + ) + + def test_freeze_model_partitions_shared_inputs(self): + # Dims and data that free RVs depend on are frozen to constants; those that only + # Deterministics and observed RVs depend on stay shared (and settable). + with pm.Model(coords={"g": [0, 1, 2], "obs": [0, 1]}) as m: + n = pm.Data("n", np.array(3, dtype="int64")) # drives a free RV size + x = pm.Data("x", np.zeros(2), dims="obs") # only feeds the observed RV + pm.Normal("z", 0, 1, shape=(n,)) + w = pm.Normal("w", 0, 1, dims="g") + pm.Normal("y", x.sum() + w.sum(), 1, observed=np.zeros(2), dims="obs") + + fm = freeze_model(m) + assert isinstance(fm["n"], Constant) + assert isinstance(fm.dim_lengths["g"], Constant) + assert isinstance(fm["x"], SharedVariable) + assert isinstance(fm.dim_lengths["obs"], SharedVariable) + assert isinstance(fm, FrozenModel) + assert not isinstance(m, FrozenModel) # original model is untouched + + def test_freeze_model_preserves_initvals(self): + # fgraph_from_model drops (and rejects) initial values; freeze_model transplants + # them onto the frozen model without touching the original. + with pm.Model() as m: + pm.Uniform("u", 0, 10, initval=7.0) # concrete + pm.Normal("p", 0, 1, size=3, initval="prior") # strategy string + + fm = freeze_model(m) + assert m.rvs_to_initial_values[m["u"]] == 7.0 # original untouched + ip = fm.initial_point(0) + np.testing.assert_allclose(ip["u_interval__"], m.initial_point(0)["u_interval__"]) + np.testing.assert_allclose(fm.initial_point(1)["p"], m.initial_point(1)["p"]) + + # Symbolic initial values reference the original graph and cannot be transplanted. + with pm.Model() as m2: + x = pm.Data("d", np.array([10.0, 20.0, 30.0])) + pm.Normal("s", 0, 1, size=3, initval=x * 2) + with pytest.raises(NotImplementedError, match="symbolic initial value"): + freeze_model(m2) + + def test_compile_fn_is_cached(self): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + mu = pm.Deterministic("mu", x * 2) + + fm = freeze_model(m) + f1 = fm.compile_fn(fm["mu"], inputs=[fm["x"]], point_fn=False) + f2 = fm.compile_fn(fm["mu"], inputs=[fm["x"]], point_fn=False) + assert f1 is f2 + assert "_compile_fn" in fm._cache + + def test_mutation_forbidden_on_frozen(self): + # FrozenModel is a sibling of Model under BaseModel: the mutating methods do not + # exist on it at all. + with pm.Model() as m: + x = pm.Uniform("x", 0, 10) + fm = freeze_model(m) + + assert isinstance(fm, FrozenModel) + assert isinstance(fm, BaseModel) + assert not isinstance(fm, pm.Model) + with pytest.raises(AttributeError, match="register_rv"): + with fm: + pm.Normal("y") + with pytest.raises(AttributeError, match="add_coord"): + fm.add_coord("new_dim", ["a", "b"]) + with pytest.raises(AttributeError, match="set_initval"): + fm.set_initval(fm["x"], 5.0) + + def test_abstract_and_final_classes(self): + with pytest.raises(TypeError, match="abstract"): + BaseModel() + with pytest.raises(TypeError, match="freeze_model"): + FrozenModel() + + def test_set_data_on_frozen_model(self): + with pm.Model() as m: + x = pm.Data("x", np.zeros(3)) + b = pm.Normal("b") + pm.Deterministic("mu", b + x) + + fm = freeze_model(m) + f1 = fm.compile_fn(fm["mu"], inputs=[fm["b"]], point_fn=False) + with fm: + pm.set_data({"x": np.ones(3)}) + f2 = fm.compile_fn(fm["mu"], inputs=[fm["b"]], point_fn=False) + assert f1 is f2 # value update reuses the cache ... + np.testing.assert_allclose(f2(0.0), np.ones(3)) # ... and flows through + + # A resize also flows through: the remaining shapes are runtime inputs. + with fm: + pm.set_data({"x": np.zeros(5)}) + f3 = fm.compile_fn(fm["mu"], inputs=[fm["b"]], point_fn=False) + assert f3 is f1 + assert f3(1.0).shape == (5,) + + def test_set_data_frozen_data_raises(self): + with pm.Model() as m: + n = pm.Data("n", np.array(3, dtype="int64")) + pm.Normal("z", 0, 1, shape=(n,)) + + fm = freeze_model(m) + with pytest.raises(TypeError, match="frozen"): + with fm: + pm.set_data({"n": np.array(5, dtype="int64")}) + + def test_compile_fn_reseeds_cached_function(self): + # A cache hit must still respect random_seed (reapplied on each compile_fn call). + with pm.Model() as m: + pm.Normal("z", 0, 1) + + fm = freeze_model(m) + f_a = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=123) + v1 = f_a() + f_b = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=123) + v2 = f_b() + assert f_a is f_b # cache hit (same compiled function reused) + np.testing.assert_allclose(v1, v2) # same seed -> same draw + + f_c = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=456) + assert not np.allclose(v1, f_c()) # different seed -> different draw + + def test_compile_fn_jax_bypasses_cache_but_stays_seeded(self): + # JAX detaches a function's RNG variables at compile time, so a cached function + # can't be reseeded. compile_fn skips the cache for RNG functions on JAX and seeds + # before compiling; deterministic functions are still cached. + pytest.importorskip("jax") + with pm.Model() as m: + pm.Normal("z", 0, 1, size=3) + b = pm.Normal("b") + pm.Deterministic("d", b * 2) + + fm = freeze_model(m) + f_a = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=0, mode="JAX") + f_b = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=0, mode="JAX") + assert f_a is not f_b # not cached (RNG function on JAX) + np.testing.assert_allclose(f_a(), f_b()) # same seed -> same draw + f_c = fm.compile_fn(fm["z"], inputs=[], point_fn=False, random_seed=1, mode="JAX") + assert not np.allclose(f_a(), f_c()) # different seed -> different draw + + g_a = fm.compile_fn(fm["d"], inputs=[fm["b"]], point_fn=False, mode="JAX") + g_b = fm.compile_fn(fm["d"], inputs=[fm["b"]], point_fn=False, mode="JAX") + assert g_a is g_b # deterministic JAX functions are cached + + def test_logp_dlogp_function_is_cached(self): + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + pm.Normal("y", x, 1, observed=[0.3, -0.5]) + + fm = freeze_model(m) + 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 + + def test_logp_dlogp_d2logp_graphs_are_cached(self): + # Memoized graph construction returns the same object, so a freshly requested logp + # graph hits the compile cache instead of recompiling. + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + pm.Normal("y", x, 1, observed=[0.3, -0.5]) + + fm = freeze_model(m) + assert fm.logp() is fm.logp() + assert fm.dlogp() is fm.dlogp() + assert fm.d2logp() is fm.d2logp() + + f1 = fm.compile_fn(fm.logp(), inputs=fm.value_vars, point_fn=False) + f2 = fm.compile_fn(fm.logp(), inputs=fm.value_vars, point_fn=False) + assert f1 is f2 + + def test_cached_logp_not_corrupted_by_dlogp(self): + # dlogp/d2logp rewrite the shared logp in place; it must stay compilable and correct. + with pm.Model() as m: + x = pm.Normal("x", 0, 1, size=2) + pm.Normal("y", x, 1, observed=[0.3, -0.5]) + fm = freeze_model(m) + _ = fm.dlogp() + _ = fm.d2logp() + f_cached = fm.compile_fn(fm.logp(), inputs=fm.value_vars, point_fn=False) + f_fresh = m.compile_fn(m.logp(), inputs=m.value_vars, point_fn=False) + + val = np.array([0.1, 0.2]) + np.testing.assert_allclose(f_cached(val), f_fresh(val)) + + def test_logp_dlogp_function_extra_values_are_per_call(self): + # A subset gradient leaves the other value var as an "extra" (shared) input. + # The cached function must reflect each call's extra values, not a stale set. + with pm.Model() as m: + x = pm.Normal("x", 0, 1) + pm.Normal("y", 0, 1) + pm.Normal("obs", m["x"] + m["y"], 1, observed=[0.0]) + + fm = freeze_model(m) + x_val = fm.rvs_to_values[fm["x"]] + f_a = fm.logp_dlogp_function([x_val], ravel_inputs=True, initial_point={"x": 0.0, "y": 1.0}) + logp_a, _ = f_a(np.array([0.0])) + f_b = fm.logp_dlogp_function([x_val], ravel_inputs=True, initial_point={"x": 0.0, "y": 2.0}) + logp_b, _ = f_b(np.array([0.0])) + assert f_a is f_b # same compiled function reused + assert not np.isclose(logp_a, logp_b) # but the extra value (y) took effect + + def test_initial_point_is_cached(self): + with pm.Model() as m: + pm.Normal("x", 0, 1, size=3) + + fm = freeze_model(m) + ip1 = fm.initial_point(0) + assert "_make_initial_point" 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_model_parent_set_programmatically(): with pm.Model() as model: x = pm.Normal("x") diff --git a/tests/model/transform/test_optimization.py b/tests/model/transform/test_optimization.py index 428bc868ce..326e9bb7c3 100644 --- a/tests/model/transform/test_optimization.py +++ b/tests/model/transform/test_optimization.py @@ -74,6 +74,26 @@ def test_freeze_dims_nothing_to_change(): assert m.point_logps() == freeze_dims_and_data(m).point_logps() +def test_freeze_dims_and_data_preserves_initvals(): + # fgraph_from_model drops (and rejects) initial values; freeze_dims_and_data preserves + # constant and strategy-string ones on the rebuilt model without touching the original. + with Model() as m: + Normal("u", initval=7.0) # concrete + Normal("p", shape=3, initval="prior") # strategy string + + frozen_m = freeze_dims_and_data(m) + assert m.rvs_to_initial_values[m["u"]] == 7.0 # original untouched + np.testing.assert_allclose(frozen_m.initial_point(0)["u"], m.initial_point(0)["u"]) + np.testing.assert_allclose(frozen_m.initial_point(1)["p"], m.initial_point(1)["p"]) + + # Symbolic initial values reference the original graph and cannot be transplanted. + with Model() as m2: + d = Data("d", [1.0, 2.0]) + Normal("s", shape=2, initval=d * 2) + with pytest.raises(NotImplementedError, match="symbolic initial value"): + freeze_dims_and_data(m2) + + def test_freeze_dims_and_data_subset(): with Model(coords={"dim1": range(3), "dim2": range(5)}) as m: data1 = Data("data1", [1, 2, 3], dims="dim1") diff --git a/tests/sampling/test_forward.py b/tests/sampling/test_forward.py index a6020616bf..92ec15c68c 100644 --- a/tests/sampling/test_forward.py +++ b/tests/sampling/test_forward.py @@ -15,6 +15,7 @@ import warnings from contextlib import nullcontext +from unittest.mock import patch import numpy as np import numpy.random as npr @@ -1389,6 +1390,77 @@ def test_pytensor_function_kwargs(self): class TestSamplePosteriorPredictive: + def test_forward_function_reused_across_set_data_batches(self): + # On a frozen model, posterior predictive sampling in a loop with set_data must + # reuse the compiled forward function instead of recompiling it on every call. + import pytensor + + from pymc.model.transform.optimization import freeze_model + + rng = np.random.default_rng(0) + N = 30 + with pm.Model() as m: + x = pm.Data("x", rng.normal(size=N)) + a = pm.Normal("a", 0, 1) + b = pm.Normal("b", 0, 1) + pm.Normal("y", a + b * x, 1, observed=rng.normal(size=N), shape=x.shape) + idata = pm.sample( + draws=20, tune=20, chains=2, progressbar=False, compute_convergence_checks=False + ) + + frozen_m = freeze_model(m) + orig_function = pytensor.function + n_compiles = [0] + + def counting_function(*args, **kwargs): + n_compiles[0] += 1 + return orig_function(*args, **kwargs) + + results = [] + with patch("pytensor.function", counting_function): + for i in range(4): + with frozen_m: + pm.set_data({"x": rng.normal(size=N)}) + pp = pm.sample_posterior_predictive(idata, progressbar=False, random_seed=i) + results.append(pp.posterior_predictive["y"].values.copy()) + # A resize is also served by the cache: the data shape is a runtime input. + with frozen_m: + pm.set_data({"x": rng.normal(size=2 * N)}) + pp = pm.sample_posterior_predictive(idata, progressbar=False, random_seed=0) + + # The forward function is compiled once and replayed for the remaining batches. + assert n_compiles[0] == 1 + # Reuse stays correct: new data flows through and draws are not frozen copies. + assert not np.allclose(results[0], results[1]) + assert pp.posterior_predictive["y"].shape[-1] == 2 * N + + def test_reused_forward_function_is_reproducible_across_seeds(self): + # Reusing the cached forward function must stay reproducible: it is reseeded per + # call, so same seed -> same draws, different seed -> different draws. + from pymc.model.transform.optimization import freeze_model + + rng = np.random.default_rng(0) + N = 20 + with pm.Model() as m: + x = pm.Data("x", rng.normal(size=N)) + a = pm.Normal("a", 0, 1) + pm.Normal("y", a * x, 1, observed=rng.normal(size=N), shape=x.shape) + idata = pm.sample( + draws=20, tune=20, chains=2, progressbar=False, compute_convergence_checks=False + ) + + with freeze_model(m): + # Same seed twice -> identical (second call reuses the cache); different seed differs. + pp_a = pm.sample_posterior_predictive(idata, progressbar=False, random_seed=42) + pp_b = pm.sample_posterior_predictive(idata, progressbar=False, random_seed=42) + pp_c = pm.sample_posterior_predictive(idata, progressbar=False, random_seed=43) + + ya = pp_a.posterior_predictive["y"].values + yb = pp_b.posterior_predictive["y"].values + yc = pp_c.posterior_predictive["y"].values + np.testing.assert_allclose(ya, yb) + assert not np.allclose(ya, yc) + def test_point_list_arg_bug_spp(self, point_list_arg_bug_fixture): pmodel, trace = point_list_arg_bug_fixture with pmodel: diff --git a/tests/test_initial_point.py b/tests/test_initial_point.py index b04542fde9..7514a02979 100644 --- a/tests/test_initial_point.py +++ b/tests/test_initial_point.py @@ -65,8 +65,8 @@ def test_dependent_initvals(self, reverse_rvs): assert ip["B1_interval__"] == 0 assert ip["B2_interval__"] == 0 - # Modify initval of L and re-evaluate - pmodel.rvs_to_initial_values[U] = 9.9 + # Modify initval of U (through the public API) and re-evaluate + pmodel.set_initval(U, 9.9) ip = pmodel.initial_point(random_seed=0) assert ip["B1_interval__"] < 0 assert ip["B2_interval__"] < 0