def gelman_rubin(chains, auto_burn_in=True):
    """ Calculates the univariate Gelman-Rubin convergence statistic
    which compares the evolution of multiple chains in a Markov-Chain Monte
    Carlo process and computes their difference to determine their convergence.
    The between-chain and within-chain variances are computed for each sampling
    parameter, and a weighted combination of the two is used to determine the
    convergence. As the chains converge, the point scale reduction factor
    should go to 1.

    Parameters
    ----------
    chains : iterable
        An iterable of numpy.array instances that contain the samples
        for each chain. Each chain has shape (nparameters, niterations).
    auto_burn_in : bool
        If True, then only use later half of samples provided.

    Returns
    -------
    psrf : numpy.array
        A numpy.array of shape (nparameters) that has the point estimates of
        the potential scale reduction factor.
    """

    # remove first half of samples
    # this will have shape (nchains, nparameters, niterations)
    if auto_burn_in:
        _, _, niterations = numpy.array(chains).shape
        chains = numpy.array([chain[:, niterations / 2 + 1:]
                              for chain in chains])

    # get number of chains, parameters, and iterations
    chains = numpy.array(chains)
    nchains, nparameters, niterations = chains.shape

    # calculate the covariance matrix for each chain
    # this will have shape (nchains, nparameters, nparameters)
    chains_covs = numpy.array([numpy.cov(chain) for chain in chains])
    if nparameters == 1:
        chains_covs = chains_covs.reshape((nchains, 1, 1))

    # calculate W the within-chain variance
    # this will have shape (nparameters, nparameters)
    w = numpy.zeros(chains_covs[0].shape)
    for i, row in enumerate(chains_covs[0]):
        for j, _ in enumerate(row):
            w[i, j] = numpy.mean(chains_covs[:, i, j])
    if nparameters == 1:
        w = w.reshape((1, 1))

    # calculate B the between-chain variance
    # this will have shape (nparameters, nparameters)
    means = numpy.zeros((nparameters, nchains))
    for i, chain in enumerate(chains):
        means[:, i] = numpy.mean(chain, axis=1).transpose()
    b = niterations * numpy.cov(means)
    if nparameters == 1:
        b = b.reshape((1, 1))

    # get diagonal elements of W and B
    # these will have shape (nparameters)
    w_diag = numpy.diag(w)
    b_diag = numpy.diag(b)

    # get variance for each chain
    # this will have shape (nparameters, nchains)
    var = numpy.zeros((nparameters, nchains))
    for i, chain_cov in enumerate(chains_covs):
        var[:, i] = numpy.diag(chain_cov)

    # get mean of means
    # this will have shape (nparameters)
    mu_hat = numpy.mean(means, axis=1)

    # get variance of variances
    # this will have shape (nparameters)
    s = numpy.var(var, axis=1)

    # get V the combined variance of all chains
    # this will have shape (nparameters)
    v = ((niterations - 1.) * w_diag / niterations +
         (1. + 1. / nchains) * b_diag / niterations)

    # get factors in variance of V calculation
    # this will have shape (nparameters)
    k = 2 * b_diag**2 / (nchains - 1)
    mid_term = numpy.cov(
        var, means**2)[nparameters:2*nparameters, 0:nparameters].T
    end_term = numpy.cov(
        var, means)[nparameters:2*nparameters, 0:nparameters].T
    wb = niterations / nchains * numpy.diag(mid_term - 2 * mu_hat * end_term)

    # get variance of V
    # this will have shape (nparameters)
    var_v = (
        (niterations - 1.) ** 2 * s +
        (1. + 1. / nchains) ** 2 * k +
        2. * (niterations - 1.) * (1. + 1. / nchains) * wb
    ) / niterations**2

    # get degrees of freedom
    # this will have shape (nparameters)
    dof = (2. * v**2) / var_v

    # more degrees of freedom factors
    # this will have shape (nparameters)
    df_adj = (dof + 3.) / (dof + 1.)

    # estimate R
    # this will have shape (nparameters)
    r2_fixed = (niterations - 1.) / niterations
    r2_random = (1. + 1. / nchains) * (1. / niterations) * (b_diag / w_diag)
    r2_estimate = r2_fixed + r2_random

    # calculate PSRF the potential scale reduction factor
    # this will have shape (nparameters)
    psrf = numpy.sqrt(r2_estimate * df_adj)

    return psrf
Calculates the univariate Gelman-Rubin convergence statistic
    which compares the evolution of multiple chains in a Markov-Chain Monte
    Carlo process and computes their difference to determine their convergence.
    The between-chain and within-chain variances are computed for each sampling
    parameter, and a weighted combination of the two is used to determine the
    convergence. As the chains converge, the point scale reduction factor
    should go to 1.

    Parameters
    ----------
    chains : iterable
        An iterable of numpy.array instances that contain the samples
        for each chain. Each chain has shape (nparameters, niterations).
    auto_burn_in : bool
        If True, then only use later half of samples provided.

    Returns
    -------
    psrf : numpy.array
        A numpy.array of shape (nparameters) that has the point estimates of
        the potential scale reduction factor.