merging changes from main

Author: michaelraczycki

pymc_experimental/model_builder.py

@@ -14,6 +14,7 @@
 
 
 import hashlib
+import json
 from abc import abstractmethod
 from pathlib import Path
 from typing import Dict, Union
@@ -188,8 +189,8 @@ def load(cls, fname: str):
         filepath = Path(str(fname))
         idata = az.from_netcdf(filepath)
         model_builder = cls(
-            dict(zip(idata.attrs["model_config_keys"], idata.attrs["model_config_values"])),
-            dict(zip(idata.attrs["sample_config_keys"], idata.attrs["sample_config_values"])),
+            json.loads(idata.attrs["model_config"]),
+            json.loads(idata.attrs["sampler_config"]),
             idata.fit_data.to_dataframe(),
         )
         model_builder.idata = idata
@@ -239,10 +240,8 @@ def fit(
         self.idata.attrs["id"] = self.id
         self.idata.attrs["model_type"] = self._model_type
         self.idata.attrs["version"] = self.version
-        self.idata.attrs["sample_config_keys"] = tuple(self.sample_config.keys())
-        self.idata.attrs["sample_config_values"] = tuple(self.sample_config.values())
-        self.idata.attrs["model_config_keys"] = tuple(self.model_config.keys())
-        self.idata.attrs["model_config_values"] = tuple(self.model_config.values())
+        self.idata.attrs["sampler_config"] = json.dumps(self.sample_config)
+        self.idata.attrs["model_config"] = json.dumps(self.model_config)
         self.idata.add_groups(fit_data=self.data.to_xarray())
         return self.idata