Skip to content

Cache model functions for iterative workflows #7815

Description

@ricardoV94

Discussed in #7177

Originally posted by ricardoV94 February 28, 2024
Users are often mislead to believe that when defining models with MutableData, PyMC will avoid recompiling logp/prior/posterior predictive functions when only the MutableData has changed. This is not the case!

PyMC always recompiles functions when you call pm.sample or pm.sample_posterior_predictive. We used to cache some of the model methods but this was removed because models are not static objects and new variables can be added any time: 983d444#diff-f593a76ecc6f9a5c5abdd7bbd3e9de9add74a068e75e64f66b7b1424a279a0dd

There's a couple approaches we could take to give back to users the caching functionality:

  1. Allow caching, but delete all cache whenever a model is mutated, which should happen only when model.register_rv is called or model.rvs_to_transforms is changed. We could easily forbid the latter now that we have https://www.pymc.io/projects/docs/en/v5.10.3/api/model/generated/pymc.model.transform.conditioning.change_value_transforms.html
  2. Provide a function to "freeze" a PyMC model. Then the model, much like nutpie compiled_pymc can never be mutated and caching would be safe.
  3. Provide a better low-level API for users to request model functions and use them in sampling routines

We may need to include the contents of pytensor.config as part of the caching / freezing so that users can request the same functions in different backends.

Linked issues:

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions