Probability: Cannot wrap a SimpleStepSizeAdaptation kernel in a TransformedTransitionKernel

[zhulingchen]

I would like to adaptively change the step_size of a HamiltonianMonteCarlo sampling kernel by wrapping HamiltonianMonteCarlo within a SimpleStepSizeAdaptation kernel and then apply some transformations to the random variables by wrapping the previous kernel into a TransformedTransitionKernel. Below is my example:

n_results = 4000
n_burnin_steps = 2000
step_size = 1e-2
num_leapfrog_steps = 10

hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=unnormalized_posterior_log_prob,
    step_size=step_size,
    num_leapfrog_steps=num_leapfrog_steps)

sssa_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
    inner_kernel=hmc_kernel, num_adaptation_steps=int(n_burnin_steps * 0.8))

tt_kernel = tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=sssa_kernel,
    bijector=unconstraining_bijectors)

[theta, mu, invcov_chol], kernel_results = tfp.mcmc.sample_chain(num_results=n_results,
                                                                 num_burnin_steps=n_burnin_steps,
                                                                 current_state=initial_state,
                                                                 kernel=tt_kernel,
                                                                 parallel_iterations=100)

But when I run it, the error is: target_log_prob_fn = inner_kernel_kwargs['target_log_prob_fn'] KeyError: 'target_log_prob_fn'

It seems that the MCMC is looking for the target_log_prob_fn key of the inner_kernel of the tt_kernel, which is sssa_kernel; but here sssa_kernel is still a "wrapped" kernel and does not have a key called target_log_prob_fn. The location of the target_log_prob_fn key is actually in the inner_kernel of sssa_kernel, i.e., hmc_kernel. It means that the "wrapping kernels" cannot be nested for now. Am I right?

Is there any workaround? Thank you!

[ElefHead]

A simple way that works is to first use hmc_kernel in tt_kernel and then wrap that in the sssa_kernel.
So, your code would change to

n_results = 4000
n_burnin_steps = 2000
step_size = 1e-2
num_leapfrog_steps = 10

hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=unnormalized_posterior_log_prob,
    step_size=step_size,
    num_leapfrog_steps=num_leapfrog_steps)

# define tt_kernel first
tt_kernel = tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=hmc_kernel,
    bijector=unconstraining_bijectors)

# defined sssa_kernel later
sssa_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
    inner_kernel=tt_kernel, # wrapped sssa kernel around tt_kernel
    num_adaptation_steps=int(n_burnin_steps * 0.8))

[theta, mu, invcov_chol], kernel_results = tfp.mcmc.sample_chain(num_results=n_results,
                                                                 num_burnin_steps=n_burnin_steps,
                                                                 current_state=initial_state,
                                                                 kernel=sssa_kernel, # changed the kernel to sssa_kernel
                                                                 parallel_iterations=100)

This should work as expected. Hope this helps, please feel free to correct me if I am wrong.