Cache model functions for iterative workflows

[ricardoV94]

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:

• Re-using models without recompiling / CompiledModel API #6239
• Allow reusing function in pymc.compute_log_likelihood #7073

[twiecki]

I like the options in this order: 1, 2, 3.

[ricardoV94]

Me too, only thing is that caching has a tendency to turn out more complex that it seems at first glance. I do think this will be less of an issue as we keep moving into a more functional workflow where we don't modify models inplace but instead recreate new ones like we do in change_value_transforms, do, observed

[zaxtax]

I'm a fan of option 1.

[twiecki]

One thing that would be nice for deployment is if I could pre-compile a model so that when a user wants to run something it doesn't have to compile at that point.

[ricardoV94]

That's a completely different challenge level, requires ahead of time compilation. I don't think any of the backends has anything good there/yet