SiegeLordEx
commented
5 years ago FWIW, a way to do it today is to write a new TransitionKernel kernel, something like this:
class Gibbs(tfp.mcmc.TransitionKernel):
def __init__(self, target_log_prob_fn, make_kernel_fns):
self._target_log_prob_fn = target_log_prob_fn
self._make_kernel_fns = make_kernel_fns
def is_calibrated(self):
return True
def one_step(self, state, _):
for i, make_kernel_fn in enumerate(self._make_kernel_fns):
def _target_log_prob_fn_part(state_part):
state[i] = state_part
return self._target_log_prob_fn(*state)
kernel = make_kernel_fn(_target_log_prob_fn_part)
state[i], _ = kernel.one_step(state[i], kernel.bootstrap_results(state[i]))
return state, ()
def bootstrap_results(self, state):
return ()Then you just pass it a list of kernel creation functions (that create a transition kernel from a target_log_prob_fn) that correspond to the elements of the state list.
Few caveats about this snippet:
- I was sloppy with respect to state parts that are lists themselves; fixing this is an exercise left to the reader.
- It won't work correctly for all types of kernels (e.g. any sort of non-Markovian kernel would be broken here).
chrism0dwk
simeoncarstens
brianwa84
Sometime when sampling from complex models, it would be handy to provide a list of kernels to tfp.mcmc.sample_chain instead of just one. As a simple example, suppose I have a joint probability distribution, with density function $\pi(\theta, \phi, \psi)$ where $-\infty < \theta, \phi < \infty$ and $\psi \in {-1, 0, 1}$. Here, I might need two kernels: an HMC for $\theta, \phi$ and a Metropolis Hastings for $\psi$.
Could we have a feature for composing such updates, as currently it is not obvious how to achieve this?
Thanks,
Chris