Closed nenasedk closed 2 months ago
Thanks for pointing this out. I made a fix for it, and will push to 2.5.2 along with some other minor improvements.
For temporary fix do this before saving:
import jax
import numpy as np
results = jax.tree.map(np.asarray, results)
Also, I notice you're using 2.5.0, due to #168 you should upgrade to >=2.5.1
Describe the bug Trying to run a toy model based on the multivariate likelihood example (https://jaxns.readthedocs.io/en/latest/examples/mvn_data_mvn_prior.html). The sampler runs fine, and I can convert the
state
into aresults
, make diagnostic plots etc. However, when running the save_results function, I get the following error:Expected behavior That the results should be saved to a .json or .npz file.
Observed behavior Full traceback:
Minimal Verifiable Complete Example
log_likelihood
andprior_model
are as in the multivariate likelihood example.JAXNS version Output of
pip freeze | grep jaxns
: jaxns==2.5.0