Compute the split-rhat for a 2d array.

def _get_split_rhat(values, round_to=2):
    """Compute the split-rhat for a 2d array."""
    shape = values.shape
    if len(shape) != 2:
        raise TypeError("Effective sample size calculation requires 2 dimensional arrays.")
    _, num_samples = shape
    num_split = num_samples // 2
    # Calculate split chain mean
    split_chain_mean1 = np.mean(values[:, :num_split], axis=1)
    split_chain_mean2 = np.mean(values[:, num_split:], axis=1)
    split_chain_mean = np.concatenate((split_chain_mean1, split_chain_mean2))
    # Calculate split chain variance
    split_chain_var1 = np.var(values[:, :num_split], axis=1, ddof=1)
    split_chain_var2 = np.var(values[:, num_split:], axis=1, ddof=1)
    split_chain_var = np.concatenate((split_chain_var1, split_chain_var2))
    # Calculate between-chain variance
    between_chain_variance = num_samples / 2 * np.var(split_chain_mean, ddof=1)
    # Calculate within-chain variance
    within_chain_variance = np.mean(split_chain_var)
    # Estimate of marginal posterior variance
    split_rhat = np.sqrt(
        (between_chain_variance / within_chain_variance + num_samples / 2 - 1) / (num_samples / 2)
    )

    return round(split_rhat, round_to)