def make_fast_adapt_kernel(*,
kind,
proposal_kernel_kwargs,
dual_averaging_kwargs):
return dassa.DualAveragingStepSizeAdaptation(
_make_base_kernel(
kind=kind, proposal_kernel_kwargs=proposal_kernel_kwargs),
**dual_averaging_kwargs)
def make_fast_adapt_kernel(*, target_log_prob_fn, initial_step_size,
num_leapfrog_steps, num_adaptation_steps,
target_accept_prob, momentum_distribution):
return dassa.DualAveragingStepSizeAdaptation(
_make_base_kernel(target_log_prob_fn=target_log_prob_fn,
step_size=initial_step_size,
num_leapfrog_steps=num_leapfrog_steps,
momentum_distribution=momentum_distribution),
shrinkage_target=initial_step_size,
num_adaptation_steps=num_adaptation_steps,
target_accept_prob=target_accept_prob)
