feature request: update function for num_leapfrog_steps

[soerenberg]

The Hamiltonian MCMC kernel tfp.mcmc.HamiltonianMonteCarlo has a very handy step_size_update_fn. I think it would be useful to have something similar for the number of leapfrog steps as well. Very often I found myself tweaking this parameter by hand, and if there was a possibility to set it even just to Floor[1/step_size] it would be very useful.

[junpenglao]

I dont think it is possible currently, as the multi-chain sampling need to synchronized across different chains, so you need to have the same number of leapfrog steps per sample. The autobatching machine @axch is working on should solve this (and NUTS).

[SiegeLordEx]

To the extent that it makes sense to share a single step size across dimensions, it makes sense to fix the product of step_size and num_leapfrog_steps the same way as well. We're working on a parameter adaptation API which will definitely make this possible.

[SiegeLordEx]

As an example of this working, see the experimental NeuTra kernel:

probability/experimental/neutra/neutra_kernel.py

Lines 369 to 381 in 8fcae25

	@tfp.mcmc.internal.util.make_innermost_setter
	def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
	return kernel_results._replace(
	accepted_results=kernel_results.accepted_results._replace(
	num_leapfrog_steps=num_leapfrog_steps))
	step_size = previous_kernel_results.new_step_size
	previous_kernel_results = set_num_leapfrog_steps(
	previous_kernel_results, self._num_leapfrog_steps(step_size))
	new_state, kernel_results = self._kernel.one_step(
	self._flatten_state(current_state), previous_kernel_results)
	return self._unflatten_state(new_state), kernel_results

If you wanted to extract that into a stable place, a possible API could look like this:

kernel = tfp.mcmc.HamiltonianMonteCarlo(...)
kernel = tfp.mcmc.ConstantTrajectoryLengthStepSizeAdaptation(kernel, trajectory_length=..., target_accept_prob=..., num_adaptation_steps=...)

Possibly with another name, and possibly using the upcoming dual-averaging step size adaptation as the implementation.