# Buggy code
def parametric(subposteriors, diagonal=False):
    """
    Merges subposteriors following (embarrassingly parallel) parametric Monte Carlo algorithm.

    **References:**

    1. *Asymptotically Exact, Embarrassingly Parallel MCMC*,
       Willie Neiswanger, Chong Wang, Eric Xing

    :param list subposteriors: a list in which each element is a collection of samples.
    :param bool diagonal: whether to compute weights using variance or covariance, defaults to
        `False` (using covariance).
    :return: the estimated mean and variance/covariance parameters of the joined posterior
    """
    joined_subposteriors = tree_multimap(lambda *args: np.stack(args), *subposteriors)
    joined_subposteriors = vmap(vmap(lambda sample: ravel_pytree(sample)[0]))(
        joined_subposteriors
    )

    submeans = np.mean(joined_subposteriors, axis=1)
    n_samples = tree_flatten(subposteriors[0])[0][0].shape[0]
    if diagonal:
        # NB: jax.numpy.var does not support ddof=1, so we do it manually
        weights = vmap(lambda x: (1 - 1 / n_samples) / np.var(x, axis=0))(
            joined_subposteriors
        )
        var = 1 / np.sum(weights, axis=0)
        normalized_weights = var * weights

        # comparing to consensus implementation, we compute weighted mean here
        mean = np.einsum("ij,ij->j", normalized_weights, submeans)
        return mean, var
    else:
        weights = vmap(lambda x: np.linalg.inv(_cov(x)))(joined_subposteriors)
        cov = np.linalg.inv(np.sum(weights, axis=0))
        normalized_weights = np.matmul(cov, weights)

        # comparing to consensus implementation, we compute weighted mean here
        mean = np.einsum("ijk,ik->j", normalized_weights, submeans)
        return mean, cov


# Traceback (RuntimeError)
test/test_examples.py::test_cpu[hmm.py --num-samples 100 --num-warmup 100 --num-chains 2] Running:
python examples/hmm.py --num-samples 100 --num-warmup 100 --num-chains 2
Simulating data...
Starting inference...
Traceback (most recent call last):
File "/home/travis/build/pyro-ppl/numpyro/examples/hmm.py", line 180, in <module>
main(args)
File "/home/travis/build/pyro-ppl/numpyro/examples/hmm.py", line 159, in main
supervised_words, unsupervised_words,
File "/home/travis/build/pyro-ppl/numpyro/numpyro/hmc_util.py", line 775, in initialize_model
raise RuntimeError("Cannot find valid initial parameters. Please check your model again.")
RuntimeError: Cannot find valid initial parameters. Please check your model again.
FAILED

# Fixed code
def parametric(subposteriors, diagonal=False):
    """
    Merges subposteriors following (embarrassingly parallel) parametric Monte Carlo algorithm.

    **References:**

    1. *Asymptotically Exact, Embarrassingly Parallel MCMC*,
       Willie Neiswanger, Chong Wang, Eric Xing

    :param list subposteriors: a list in which each element is a collection of samples.
    :param bool diagonal: whether to compute weights using variance or covariance, defaults to
        `False` (using covariance).
    :return: the estimated mean and variance/covariance parameters of the joined posterior
    """
    joined_subposteriors = tree_multimap(lambda *args: np.stack(args), *subposteriors)
    joined_subposteriors = vmap(vmap(lambda sample: ravel_pytree(sample)[0]))(
        joined_subposteriors
    )

    submeans = np.mean(joined_subposteriors, axis=1)
    if diagonal:
        # NB: jax.numpy.var does not support ddof=1, so we do it manually
        weights = vmap(lambda x: 1 / np.var(x, ddof=1, axis=0))(joined_subposteriors)
        var = 1 / np.sum(weights, axis=0)
        normalized_weights = var * weights

        # comparing to consensus implementation, we compute weighted mean here
        mean = np.einsum("ij,ij->j", normalized_weights, submeans)
        return mean, var
    else:
        weights = vmap(lambda x: np.linalg.inv(np.cov(x.T)))(joined_subposteriors)
        cov = np.linalg.inv(np.sum(weights, axis=0))
        normalized_weights = np.matmul(cov, weights)

        # comparing to consensus implementation, we compute weighted mean here
        mean = np.einsum("ijk,ik->j", normalized_weights, submeans)
        return mean, cov