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. Default value: None.
  • 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).