code: def get_sample_stats(fit, log_likelihood=None):
      
      dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64}
  
      ndraws = [s - w for s, w in zip(fit.sim["n_save"], fit.sim["warmup2"])]
  
      extraction = OrderedDict()
      for chain, (pyholder, ndraws) in enumerate(zip(fit.sim["samples"], ndraws)):
          if chain == 0:
              for key in pyholder["sampler_param_names"]:
                  extraction[key] = []
          for key, values in zip(pyholder["sampler_param_names"], pyholder["sampler_params"]):
              extraction[key].append(values[-ndraws:])
  
      data = OrderedDict()
      for key, values in extraction.items():
          values = np.stack(values, axis=0)
          dtype = dtypes.get(key)
          values = values.astype(dtype)
          name = re.sub("__$", "", key)
          name = "diverging" if name == "divergent" else name
          data[name] = values
  
      # log_likelihood
      if log_likelihood is not None:
          log_likelihood_data = get_draws(fit, variables=log_likelihood)
          data["log_likelihood"] = log_likelihood_data[log_likelihood]
  
      # lp__
      stat_lp = get_draws(fit, variables="lp__")
      data["lp"] = stat_lp["lp__"]
  
      return data
docs: Extract sample stats from PyStan fit.