.wip force load

Author: ricardoV94

pymc_experimental/tests/utils/test_cache.py

@@ -17,8 +17,13 @@ def test_cache_sampling(tmpdir):
         assert len(os.listdir(tmpdir)) == 0
 
         prior1, prior2 = (cache_prior(samples=5) for _ in range(2))
+        prior3 = cache_sampling(pm.sample_prior_predictive, dir=tmpdir, force_sample=True)(
+            samples=5
+        )
         assert len(os.listdir(tmpdir)) == 1
         assert prior1.prior["x"].mean() == prior2.prior["x"].mean()
+        assert prior2.prior["x"].mean() != prior3.prior["x"].mean()
+        assert prior2.prior_predictive["y"].mean() != prior3.prior_predictive["y"].mean()
 
         post1, post2 = (cache_post(tune=5, draws=5, progressbar=False) for _ in range(2))
         assert len(os.listdir(tmpdir)) == 2

pymc_experimental/utils/cache.py

@@ -76,6 +76,7 @@ def cache_sampling(
     sampling_fn: Literal[sample, sample_prior_predictive, sample_posterior_predictive],
     dir: str = "",
     force_sample: bool = False,
+    force_load: bool = True,
 ) -> Callable:
     """Cache the result of PyMC sampling.
 
@@ -88,6 +89,7 @@ def cache_sampling(
         The directory where the results should be saved or retrieved from. Defaults to working directory.
     force_sample: bool, Optional
         Whether to force sampling even if cache is found. Defaults to False.
+    force_load:
 
     Returns
     -------
@@ -163,7 +165,9 @@ def wrapped_sampling_fn(*args, model=None, random_seed=None, **kwargs):
 
         if not force_sample and os.path.exists(file_path):
             print("Cache hit! Returning stored result", file=sys.stdout)
-            idata_out = az.from_netcdf(file_path)
+            idata_out: az.InferenceData = az.from_netcdf(file_path)
+            if force_load:
+                idata_out.load()
 
         else:
             idata_out = sampling_fn(*args, **kwargs, model=model, random_seed=random_seed)