Add function that caches sampling results

[ricardoV94]

import pymc as pm
from pymc_experimental.utils.cache import cache_sampling

with pm.Model() as m:
    y_data = pm.MutableData("y_data", [0, 1, 2])
    x = pm.Normal("x", 0, 1)
    y = pm.Normal("y", mu=x, observed=y_data)

    cache_sample = cache_sampling(pm.sample, dir="traces")
    idata1 = cache_sample(chains=2)

    # Cache hit! Returning stored result
    idata2 = cache_sample(chains=2)

    pm.set_data({"y_data": [1, 1, 1]})
    idata3 = cache_sample(chains=2)

assert idata1.posterior["x"].mean() == idata2.posterior["x"].mean()
assert idata1.posterior["x"].mean() != idata3.posterior["x"].mean()

[twiecki]

When would that be useful?

[ricardoV94]

When rerunning notebooks or any workflow with saving/loading of traces where you might still be tinkering with the model.

You don't need to bother defining the names of the traces, or overriding old traces, since caching is automatically derived from the model and its data

[ferrine]

maybe to pytensor?

[ricardoV94]

Too experimental for now

[fonnesbeck]

I usually rely on things like MLFlow for storing artifacts like this.

[ricardoV94]

I'm not familiar with MLflow, the idea here is that it pairs the saved traces to the exact model/sampling function (and arguments) that were used.

Basically the model and the function kwargs are the cache key.

Does this have any parallel to your workflow with MLflow?

[lucianopaz]

Suggested change

	name = str(obj)
	props = str(getattr(obj, "_props", lambda: {})())
	name = str(obj)
	if hasattr(obj, "_props"):
	prop_dict = obj._prop_dict()
	props = str(
	{k: get_name_and_props(v) for k, v in prop_dict.items()}
	)
	else:
	props = str({})
	name = str(obj)
	props = str(getattr(obj, "_props", lambda: {})())

[lucianopaz]

This is just to make sure that potential recursion of _props is handled.

[ricardoV94]

_props are not recursive

[lucianopaz]

Suggested change

	os.mkdir(dir)
	os.makedirs(dir, exist_ok=True)

I think that it's better to use os.makedirs because it creates the intermediate directories if they are required.

[lucianopaz]

This looks weird to me. The sampling_fn will go make up the hash. So if someone first calls sample, and then sample_posterior_predictive using extend, the second cache will include both the posterior and posterior_predictive groups; but the first cache will only include the posterior group. I think that it's cleaner to never allow for extending the idata inplace, and force users to combine the different InferenceData objects themselves.

[ricardoV94]

Good catch

[ricardoV94]

Wait no, I don't see what the problem is. We only ever save new idatas coming out of the sampling_fn and these never extend the previous one.