Cache compiled model functions

conda-envs/environment-alternative-backends.yml

@@ -13,13 +13,11 @@ dependencies:
 - numba
 # nutpie is installed from git in the CI workflow (see tests.yml) so we can track
 # the arviz 1.0 -compatible branch until an upstream release ships.
-# Jaxlib version must not be greater than jax version!
-- jax>=0.4.28
-- jaxlib>=0.4.28
+# The JAX stack (jax, jaxlib, numpyro, blackjax) is installed from PyPI in the pip
+# section below rather than conda-forge -- see the note there.
 - libblas=*=*mkl
 - mkl-service
 - numpy>=1.25.0
-- numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
 - pytensor>=3.1.2,<3.2
@@ -39,7 +37,15 @@ dependencies:
 - types-cachetools
 # blackjax now does not pull in fastprogress by default
 - fastprogress>=1.0.0
-- blackjax>=1.5
 - pip:
   - numdifftools>=0.9.40
   - mcbackend>=0.4.0
+  # conda-forge builds jaxlib from source with its own XLA/LLVM toolchain, and its
+  # 0.10.x CPU build segfaults in XLA codegen while compiling some graphs
+  # (DirichletMultinomial, MvStudentT draws) on the CI runners. The official PyPI
+  # wheels compile the same graphs fine, so install the whole JAX stack from PyPI.
+  # Jaxlib version must not be greater than jax version!
+  - jax>=0.4.28
+  - jaxlib>=0.4.28
+  - numpyro>=0.8.0
+  - blackjax>=1.5

docs/source/api/model/core.rst

@@ -5,6 +5,8 @@ Model creation and inspection
 .. autosummary::
    :toctree: generated/
 
+   BaseModel
+   FrozenModel
    Model
    modelcontext
 

docs/source/api/model/optimization.rst

@@ -5,3 +5,4 @@ Model Optimization
    :toctree: generated/
 
    freeze_dims_and_data
+   freeze_model

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.

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,
     ):

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.

pymc/dims/model.py

@@ -23,20 +23,22 @@
     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.
 
     Dimensions are required if the input is not a scalar.
     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)