Extract sample_stats from fit.

def sample_stats_to_xarray(self):
        """Extract sample_stats from fit."""
        dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64}

        # copy dims and coords
        dims = deepcopy(self.dims) if self.dims is not None else {}
        coords = deepcopy(self.coords) if self.coords is not None else {}

        sampler_params = self.sample_stats
        log_likelihood = self.log_likelihood
        if isinstance(log_likelihood, str):
            log_likelihood_cols = [
                col for col in self.posterior[0].columns if log_likelihood == col.split(".")[0]
            ]
            log_likelihood_vals = [item[log_likelihood_cols] for item in self.posterior]

            # Add log_likelihood to sampler_params
            for i, _ in enumerate(sampler_params):
                # slice log_likelihood to keep dimensions
                for col in log_likelihood_cols:
                    col_ll = col.replace(log_likelihood, "log_likelihood")
                    sampler_params[i][col_ll] = log_likelihood_vals[i][col]

            # change dims and coords for log_likelihood if defined
            if log_likelihood in dims:
                dims["log_likelihood"] = dims.pop(log_likelihood)

            log_likelihood_dims = np.array(
                [list(map(int, col.split(".")[1:])) for col in log_likelihood_cols]
            )
            max_dims = log_likelihood_dims.max(0)
            max_dims = max_dims if hasattr(max_dims, "__iter__") else (max_dims,)
            default_dim_names, _ = generate_dims_coords(shape=max_dims, var_name=log_likelihood)
            log_likelihood_dim_names, _ = generate_dims_coords(
                shape=max_dims, var_name="log_likelihood"
            )
            for default_dim_name, log_likelihood_dim_name in zip(
                default_dim_names, log_likelihood_dim_names
            ):
                if default_dim_name in coords:
                    coords[log_likelihood_dim_name] = coords.pop(default_dim_name)

        for j, s_params in enumerate(sampler_params):
            rename_dict = {}
            for key in s_params:
                key_, *end = key.split(".")
                name = re.sub("__$", "", key_)
                name = "diverging" if name == "divergent" else name
                rename_dict[key] = ".".join((name, *end))
                sampler_params[j][key] = s_params[key].astype(dtypes.get(key))
            sampler_params[j] = sampler_params[j].rename(columns=rename_dict)
        data = _unpack_dataframes(sampler_params)
        return dict_to_dataset(data, coords=coords, dims=dims)