code: def rhat(data, var_names=None):
      r
      if isinstance(data, np.ndarray):
          return _get_split_rhat(data)
  
      dataset = convert_to_dataset(data, group="posterior")
      var_names = _var_names(var_names, dataset)
  
      dataset = dataset if var_names is None else dataset[var_names]
      return xr.apply_ufunc(_rhat_ufunc, dataset, input_core_dims=(("chain", "draw"),))
docs: r"""Compute estimate of Split R-hat for a set of traces.
  
      The Split R-hat diagnostic tests for lack of convergence by comparing the variance between
      multiple chains to the variance within each chain. If convergence has been achieved, the
      between-chain and within-chain variances should be identical. To be most effective in
      detecting evidence for nonconvergence, each chain should have been initialized to starting
      values that are dispersed relative to the target distribution.
  
      Parameters
      ----------
      data : obj
          Any object that can be converted to an az.InferenceData object
          Refer to documentation of az.convert_to_dataset for details
          At least 2 posterior chains are needed to compute this diagnostic of one or more
          stochastic parameters.
          For ndarray: shape = (chain, draw).
      var_names : list
        Names of variables to include in the rhat report
  
      Returns
      -------
      r_hat : xarray.Dataset
        Returns dictionary of the potential scale reduction factors, :math:`\hat{R}`
  
      Notes
      -----
      The diagnostic is computed by:
  
        .. math:: \hat{R} = \frac{\hat{V}}{W}
  
      where :math:`W` is the within-chain variance and :math:`\hat{V}` is the posterior variance
      estimate for the pooled traces. This is the potential scale reduction factor, which converges
      to unity when each of the traces is a sample from the target posterior. Values greater than one
      indicate that one or more chains have not yet converged.
  
      References
      ----------
      Gelman et al. BDA (2014)
      Brooks and Gelman (1998)
      Gelman and Rubin (1992)