JAX implementation of emcee

Greetings!

I've ported a subset of emcee functionality to the NumPyro project under the sampler name AIES.

(For the uninitiated, NumPyro uses JAX, a library with an interface to numpy and additional features like JIT compiling and GPU support, in the backend. The upshot is that if you're using currently using emcee, switching to NumPyro may give you a dramatic inference speedup!)

I've tried my best to match the existing API. You can use either the NumPyro model specification language

import jax
import jax.numpy as jnp

import numpyro
from numpyro.infer import MCMC, AIES
import numpyro.distributions as dist

n_dim, num_chains = 5, 100
mu, sigma = jnp.zeros(n_dim), jnp.ones(n_dim)

def model(mu, sigma):
    with numpyro.plate('n_dim', n_dim):
        numpyro.sample("x", dist.Normal(mu, sigma))

kernel = AIES(model, moves={AIES.DEMove() : 0.5,
                            AIES.StretchMove() : 0.5})

mcmc = MCMC(kernel, 
            num_warmup=1000,
            num_samples=2000, 
            num_chains=num_chains, 
            chain_method='vectorized')

mcmc.run(jax.random.PRNGKey(0), mu, sigma)
mcmc.print_summary()

or provide your own potential function.

def potential_fn(z):
    return 0.5 * jnp.sum(((z - mu) / sigma) ** 2)

kernel = AIES(potential_fn=potential_fn, 
              moves={AIES.DEMove() : 0.5,
                     AIES.StretchMove() : 0.5})
mcmc = MCMC(kernel, 
            num_warmup=1000,
            num_samples=2000, 
            num_chains=num_chains, 
            chain_method='vectorized')

init_params = jax.random.normal(jax.random.PRNGKey(0), 
                                (num_chains, n_dim))

mcmc.run(jax.random.PRNGKey(1), mu, sigma, init_params=init_params)
mcmc.print_summary()

Hope this is helpful to some folks!

No emcee file, test, or entry point is identified; the issue mainly presents an external NumPyro AIES implementation. Start by reviewing the linked NumPyro AIES documentation and determine whether the project wants a concrete integration or other change, since no completion criteria are specified.