arviz.from_pyjags() ValueError: too many values to unpack (expected 3)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/home/anaconda/workspace/group_code/long_rt/simulation1/jags_test.ipynb Cell 12' in <cell line: 1>()
----> 1 az.from_pyjags(tr)

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:374, in from_pyjags(posterior, prior, log_likelihood, coords, dims, save_warmup, warmup_iterations)
    313 def from_pyjags(
    314     posterior: tp.Optional[tp.Mapping[str, np.ndarray]] = None,
    315     prior: tp.Optional[tp.Mapping[str, np.ndarray]] = None,
   (...)
    320     warmup_iterations: int = 0,
    321 ) -> InferenceData:
    322     """
    323     Convert PyJAGS posterior samples to an ArviZ inference data object.
    324 
   (...)
    364     InferenceData
    365     """
    366     return PyJAGSConverter(
    367         posterior=posterior,
    368         prior=prior,
    369         log_likelihood=log_likelihood,
    370         dims=dims,
    371         coords=coords,
    372         save_warmup=save_warmup,
    373         warmup_iterations=warmup_iterations,
--> 374     ).to_inference_data()

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:107, in PyJAGSConverter.to_inference_data(self)
    103 save_warmup = self.save_warmup and self.warmup_iterations > 0
    104 # self.posterior is not None
    106 idata_dict = {
--> 107     "posterior": self.posterior_to_xarray(),
    108     "prior": self.prior_to_xarray(),
    109     "log_likelihood": self.log_likelihood_to_xarray(),
    110     "save_warmup": save_warmup,
    111 }
    113 return InferenceData(**idata_dict)

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:83, in PyJAGSConverter.posterior_to_xarray(self)
     80 if self.posterior is None:
     81     return None
---> 83 return self._pyjags_samples_to_xarray(self.posterior)

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:62, in PyJAGSConverter._pyjags_samples_to_xarray(self, pyjags_samples)
     59 def _pyjags_samples_to_xarray(
     60     self, pyjags_samples: tp.Mapping[str, np.ndarray]
     61 ) -> tp.Tuple[xarray.Dataset, xarray.Dataset]:
---> 62     data, data_warmup = get_draws(
     63         pyjags_samples=pyjags_samples,
     64         warmup_iterations=self.warmup_iterations,
     65         warmup=self.save_warmup,
     66     )
     68     return (
     69         dict_to_dataset(data, library=self.pyjags, coords=self.coords, dims=self.dims),
     70         dict_to_dataset(
   (...)
     75         ),
     76     )

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:165, in get_draws(pyjags_samples, variables, warmup, warmup_iterations)
    161         data_warmup = _convert_pyjags_dict_to_arviz_dict(
    162             samples=warmup_samples, variable_names=variables
    163         )
    164 else:
--> 165     data = _convert_pyjags_dict_to_arviz_dict(samples=pyjags_samples, variable_names=variables)
    167 return data, data_warmup

File ~/anaconda3/envs/mcmc/lib/python3.10/site-packages/arviz/data/io_pyjags.py:238, in _convert_pyjags_dict_to_arviz_dict(samples, variable_names)
    236 for variable_name, chains in samples.items():
    237     if variable_name in variable_names:
--> 238         parameter_dimension, _, _ = chains.shape
    239         if parameter_dimension == 1:
    240             variable_name_to_samples_map[variable_name] = chains[0, :, :].transpose()

ValueError: too many values to unpack (expected 3)

将源码io_pyjags.py的236到244行

    for variable_name, chains in samples.items():
        if variable_name in variable_names:
            parameter_dimension, _, _ = chains.shape
            if parameter_dimension == 1:
                variable_name_to_samples_map[variable_name] = chains[0, :, :].transpose()
            else:
                variable_name_to_samples_map[variable_name] = np.swapaxes(chains, 0, 2)

    return variable_name_to_samples_map

替换为

    for variable_name, chains in samples.items():
        if variable_name in variable_names:
            # parameter_dimension, _, _ = chain_shape
            if (len(chains.shape) == 3) and (chains.shape[2] == 1):
                variable_name_to_samples_map[variable_name] = chains[0, :, :].transpose()
            else:
                chain_axis = np.arange(len(chains.shape))
                map_axis = np.concatenate([chain_axis[-2:][::-1],
                                           chain_axis[:-2]])
                variable_name_to_samples_map[variable_name] = chains.transpose(map_axis)
    return variable_name_to_samples_map