Implementing Delayed Rejection Adaptive Metropolis (DRAM)

rahul · September 25, 2023, 3:49pm

Hello everyone,

I am interested in implementing the Delayed Rejection Adaptive Metropolis (DRAM) algorithm in pyro.

I was wondering what is the best way to go about doing this. The algorithm is similar to Metropolis Hastings, except that there are some local and global changes to the proposal distribution. Are there any MH kernels that have been implemented in pyro that I can try and use as a starting point to implement this?

I found this one, but it uses numpyro.

gist.github.com

https://gist.github.com/aflaxman/8d52ded24d0279b1dbbe74dd91c283b6

numpyro_metropolis_kernel_2.py

from collections import namedtuple
from jax import random
from jax.flatten_util import ravel_pytree
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, util
from numpyro.infer.mcmc import MCMCKernel

MetState = namedtuple("MetState", ["z", "rng_key"])

This file has been truncated. show original