Stay organized with collections Save and categorize content based on your preferences.

Create a function implementing a step-size update policy.

The simple policy increases or decreases the step_size_var based on the average of exp(minimum(0., log_accept_ratio)). It is based on Section 4.2 of Andrieu and Thoms (2008).

The num_adaptation_steps argument is set independently of any burnin for the overall chain. In general, adaptation prevents the chain from reaching a stationary distribution, so obtaining consistent samples requires num_adaptation_steps be set to a value somewhat smaller than the number of burnin steps. However, it may sometimes be helpful to set num_adaptation_steps to a larger value during development in order to inspect the behavior of the chain during adaptation.

num_adaptation_steps Scalar int Tensor number of initial steps to during which to adjust the step size. This may be greater, less than, or equal to the number of burnin steps. If None, the step size is adapted on every step (note this breaks stationarity of the chain!).
target_rate Scalar Tensor representing desired accept_ratio. Default value: 0.75 (i.e., center of asymptotically optimal rate).
decrement_multiplier Tensor representing amount to downscale current step_size. Default value: 0.01.
increment_multiplier Tensor representing amount to upscale current step_size. Default value: 0.01.
step_counter Scalar int Variable specifying the current step. The step size is adapted iff step_counter < num_adaptation_steps. Default value: if None, an internal variable step_size_adaptation_step_counter is created and initialized to -1.

step_size_simple_update_fn Callable that takes args step_size_var, kernel_results and returns updated step size(s).