class_docstr: Adapts the inner kernel's `step_size` based on `log_accept_prob`.
  
  The dual averaging policy uses a noisy step size for exploration, while
  averaging over tuning steps to provide a smoothed estimate of an optimal
  value. It is based on [section 3.2 of Hoffman and Gelman (2013)][1], which
  modifies the [stochastic convex optimization scheme of Nesterov (2009)][2].
  The modified algorithm applies extra weight to recent iterations while
  keeping the convergence guarantees of Robbins-Monro, and takes care not
  to make the step size too small too quickly when maintaining a constant
  trajectory length, to avoid expensive early iterations. A good target
  acceptance probability depends on the inner kernel. If this kernel is
  `HamiltonianMonteCarlo`, then 0.6-0.9 is a good range to aim for. For
  `RandomWalkMetropolis` this should be closer to 0.25. See the individual
  kernels' docstrings for guidance.
  
  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][3] 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.
  
  The step size is assumed to broadcast with the chain state, potentially having
  leading dimensions corresponding to multiple chains. When there are fewer of
  those leading dimensions than there are chain dimensions, the corresponding
  dimensions in the `log_accept_prob` are averaged (in the direct space, rather
  than the log space) before being used to adjust the step size. This means that
  this kernel can do both cross-chain adaptation, or per-chain step size
  adaptation, depending on the shape of the step size.
  
  For example, if your problem has a state with shape `[S]`, your chain state
  has shape `[C0, C1, S]` (meaning that there are `C0 * C1` total chains) and
  `log_accept_prob` has shape `[C0, C1]` (one acceptance probability per chain),
  then depending on the shape of the step size, the following will happen:
  
  - Step size has shape [], [S] or [1], the `log_accept_prob` will be averaged
    across its `C0` and `C1` dimensions. This means that you will learn a shared
    step size based on the mean acceptance probability across all chains. This
    can be useful if you don't have a lot of steps to adapt and want to average
    away the noise.
  
  - Step size has shape [C1, 1] or [C1, S], the `log_accept_prob` will be
    averaged across its `C0` dimension. This means that you will learn a shared
    step size based on the mean acceptance probability across chains that share
    the coordinate across the `C1` dimension. This can be useful when the `C1`
    dimension indexes different distributions, while `C0` indexes replicas of a
    single distribution, all sampled in parallel.
  
  - Step size has shape [C0, C1, 1] or [C0, C1, S], then no averaging will
    happen. This means that each chain will learn its own step size. This can be
    useful when all chains are sampling from different distributions. Even when
    all chains are for the same distribution, this can help during the initial
    warmup period.
  
  - Step size has shape [C0, 1, 1] or [C0, 1, S], the `log_accept_prob` will be
    averaged across its `C1` dimension. This means that you will learn a shared
    step size based on the mean acceptance probability across chains that share
    the coordinate across the `C0` dimension. This can be useful when the `C0`
    dimension indexes different distributions, while `C1` indexes replicas of a
    single distribution, all sampled in parallel.
  
  By default, the averaging function used above is the arithmetic mean, which is
  not robust to stuck chains (e.g. average of one chain with `p_accept = 0` and
  three chains with `p_accept = 1` will result in an average `p_accept = 0.75`,
  which will cause this kernel keep the step size roughly the same rather than
  reducing it to unstick the stuck chain). A more robust choice would be to set
  `reduce_fn` argument to `tfp.math.reduce_log_harmonic_mean_exp` [4]. Note,
  however, that the harmonic mean of a set of numbers is usually smaller than
  the arithmetic mean, so its use will typically produce smaller than optimal
  step sizes even for well behaved target distributions.
  
  #### Examples
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions
  
  target_log_prob_fn = tfd.Normal(loc=0., scale=1.).log_prob
  num_burnin_steps = 500
  num_results = 500
  num_chains = 64
  step_size = 0.1
  # Or, if you want per-chain step size:
  # step_size = tf.fill([num_chains], step_size)
  
  kernel = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=target_log_prob_fn,
      num_leapfrog_steps=2,
      step_size=step_size)
  kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
      inner_kernel=kernel, num_adaptation_steps=int(num_burnin_steps * 0.8))
  
  # The chain will be stepped for num_results + num_burnin_steps, adapting for
  # the first num_adaptation_steps.
  samples, [step_size, log_accept_ratio] = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=tf.zeros(num_chains),
      kernel=kernel,
      trace_fn=lambda _, pkr: [pkr.inner_results.accepted_results.step_size,
                               pkr.inner_results.log_accept_ratio])
  
  # ~0.75
  p_accept = tf.math.exp(tfp.math.reduce_logmeanexp(tf.minimum(
      log_accept_ratio, 0.)))
  ```
  #### References
  
  [1]: Matthew D. Hoffman, Andrew Gelman. The No-U-Turn Sampler: Adaptively
      Setting Path Lengths in Hamiltonian Monte Carlo.
      In _Journal of Machine Learning Research_, 15(1):1593-1623, 2014.
      http://jmlr.org/papers/volume15/hoffman14a/hoffman14a.pdf
  
  [2] Yurii Nesterov. Primal-dual subgradient methods for convex problems.
      Mathematical programming 120.1 (2009): 221-259
      https://link.springer.com/article/10.1007/s10107-007-0149-x
  
  [3]:
  http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745
  
  [4]: Hoffman, M., Radul, A., & Sountsov, P. An Adaptive MCMC Scheme
       for Setting Trajectory Lengths in Hamiltonian Monte Carlo, 2020. In
       preparation.
docstr: Initializes this transition kernel.
  
  Args:
    inner_kernel: `TransitionKernel`-like object.
    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.
    target_accept_prob: A floating point `Tensor` representing desired
      acceptance probability. Must be a positive number less than 1. This can
      either be a scalar, or have shape [num_chains]. Default value: `0.75`
        (the [center of asymptotically optimal rate for HMC][1]).
    exploration_shrinkage: Floating point scalar `Tensor`. How strongly the
      exploration rate is biased towards the shrinkage target.
    shrinkage_target: `Tensor` or list of tensors. Value the exploration
      step size(s) is/are biased towards.
      As `num_adaptation_steps --> infinity`, this bias goes to zero.
      Defaults to 10 times the initial step size.
    step_count_smoothing: Int32 scalar `Tensor`. Number of "pseudo-steps"
      added to the number of steps taken to prevents noisy exploration during
      the early samples.
    decay_rate: Floating point scalar `Tensor`. How much to favor recent
      iterations over earlier ones. A value of 1 gives equal weight to all
      history. A value of 0 gives weight only to the most recent iteration.
    step_size_setter_fn: A callable with the signature `(kernel_results,
      new_step_size) -> new_kernel_results` where `kernel_results` are the
      results of the `inner_kernel`, `new_step_size` is a `Tensor` or a nested
      collection of `Tensor`s with the same structure as returned by the
      `step_size_getter_fn`, and `new_kernel_results` are a copy of
      `kernel_results` with the step size(s) set.
    step_size_getter_fn: A callable with the signature `(kernel_results) ->
      step_size` where `kernel_results` are the results of the `inner_kernel`,
      and `step_size` is a floating point `Tensor` or a nested collection of
      such `Tensor`s.
    log_accept_prob_getter_fn: A callable with the signature `(kernel_results)
      -> log_accept_prob` where `kernel_results` are the results of the
      `inner_kernel`, and `log_accept_prob` is a floating point `Tensor`.
      `log_accept_prob` can either be a scalar, or have shape [num_chains]. If
      it's the latter, `step_size` should also have the same leading
      dimension.
    reduce_fn: A callable with signature `(input_tensor, axis, keepdims) ->
      tensor` that returns a log-reduction of `log_accept_prob`, typically
      some sort of mean. By default, this performs an arithmetic mean.
    experimental_reduce_chain_axis_names: A `str` or list of `str`s indicating
      the named axes that should additionally reduced during the log-reduction
      of `log_accept_prob`.
    validate_args: Python `bool`. When `True` kernel parameters are checked
      for validity. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'dual_averaging_step_size_adaptation').
docstr_tok: ["Initializes","this","transition","kernel",".","Args",":","inner_kernel",":","`","TransitionKernel","`","-like","object",".","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",".","target_accept_prob",":","A","floating","point","`","Tensor","`","representing","desired","acceptance","probability",".","Must","be","a","positive","number","less","than","1",".","This","can","either","be","a","scalar",",","or","have","shape","[","num_chains","]",".","Default","value",":","`","0.75","`","(","the","[","center","of","asymptotically","optimal","rate","for","HMC","]","[","1","]",")",".","exploration_shrinkage",":","Floating","point","scalar","`","Tensor","`",".","How","strongly","the","exploration","rate","is","biased","towards","the","shrinkage","target",".","shrinkage_target",":","`","Tensor","`","or","list","of","tensors",".","Value","the","exploration","step","size","(","s",")","is\/are","biased","towards",".","As","`","num_adaptation_steps","--",">","infinity","`",",","this","bias","goes","to","zero",".","Defaults","to","10","times","the","initial","step","size",".","step_count_smoothing",":","Int32","scalar","`","Tensor","`",".","Number","of","``","pseudo-steps","''","added","to","the","number","of","steps","taken","to","prevents","noisy","exploration","during","the","early","samples",".","decay_rate",":","Floating","point","scalar","`","Tensor","`",".","How","much","to","favor","recent","iterations","over","earlier","ones",".","A","value","of","1","gives","equal","weight","to","all","history",".","A","value","of","0","gives","weight","only","to","the","most","recent","iteration",".","step_size_setter_fn",":","A","callable","with","the","signature","`","(","kernel_results",",","new_step_size",")","-",">","new_kernel_results","`","where","`","kernel_results","`","are","the","results","of","the","`","inner_kernel","`",",","`","new_step_size","`","is","a","`","Tensor","`","or","a","nested","collection","of","`","Tensor","`","s","with","the","same","structure","as","returned","by","the","`","step_size_getter_fn","`",",","and","`","new_kernel_results","`","are","a","copy","of","`","kernel_results","`","with","the","step","size","(","s",")","set",".","step_size_getter_fn",":","A","callable","with","the","signature","`","(","kernel_results",")","-",">","step_size","`","where","`","kernel_results","`","are","the","results","of","the","`","inner_kernel","`",",","and","`","step_size","`","is","a","floating","point","`","Tensor","`","or","a","nested","collection","of","such","`","Tensor","`","s",".","log_accept_prob_getter_fn",":","A","callable","with","the","signature","`","(","kernel_results",")","-",">","log_accept_prob","`","where","`","kernel_results","`","are","the","results","of","the","`","inner_kernel","`",",","and","`","log_accept_prob","`","is","a","floating","point","`","Tensor","`",".","`","log_accept_prob","`","can","either","be","a","scalar",",","or","have","shape","[","num_chains","]",".","If","it","'s","the","latter",",","`","step_size","`","should","also","have","the","same","leading","dimension",".","reduce_fn",":","A","callable","with","signature","`","(","input_tensor",",","axis",",","keepdims",")","-",">","tensor","`","that","returns","a","log-reduction","of","`","log_accept_prob","`",",","typically","some","sort","of","mean",".","By","default",",","this","performs","an","arithmetic","mean",".","experimental_reduce_chain_axis_names",":","A","`","str","`","or","list","of","`","str","`","s","indicating","the","named","axes","that","should","additionally","reduced","during","the","log-reduction","of","`","log_accept_prob","`",".","validate_args",":","Python","`","bool","`",".","When","`","True","`","kernel","parameters","are","checked","for","validity",".","When","`","False","`","invalid","inputs","may","silently","render","incorrect","outputs",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","function",".","Default","value",":","`","None","`","(","i.e.",",","'dual_averaging_step_size_adaptation","'",")","."]
code: def __init__(
      self,
      inner_kernel,
      num_adaptation_steps,
      target_accept_prob=0.75,
      exploration_shrinkage=0.05,
      shrinkage_target=None,
      step_count_smoothing=10,
      decay_rate=0.75,
      step_size_setter_fn=hmc_like_step_size_setter_fn,
      step_size_getter_fn=hmc_like_step_size_getter_fn,
      log_accept_prob_getter_fn=hmc_like_log_accept_prob_getter_fn,
      reduce_fn=reduce_logmeanexp,
      experimental_reduce_chain_axis_names=None,
      validate_args=False,
      name=None,
  ):
      """Initializes this transition kernel.
  
      Args:
        inner_kernel: `TransitionKernel`-like object.
        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.
        target_accept_prob: A floating point `Tensor` representing desired
          acceptance probability. Must be a positive number less than 1. This can
          either be a scalar, or have shape [num_chains]. Default value: `0.75`
            (the [center of asymptotically optimal rate for HMC][1]).
        exploration_shrinkage: Floating point scalar `Tensor`. How strongly the
          exploration rate is biased towards the shrinkage target.
        shrinkage_target: `Tensor` or list of tensors. Value the exploration
          step size(s) is/are biased towards.
          As `num_adaptation_steps --> infinity`, this bias goes to zero.
          Defaults to 10 times the initial step size.
        step_count_smoothing: Int32 scalar `Tensor`. Number of "pseudo-steps"
          added to the number of steps taken to prevents noisy exploration during
          the early samples.
        decay_rate: Floating point scalar `Tensor`. How much to favor recent
          iterations over earlier ones. A value of 1 gives equal weight to all
          history. A value of 0 gives weight only to the most recent iteration.
        step_size_setter_fn: A callable with the signature `(kernel_results,
          new_step_size) -> new_kernel_results` where `kernel_results` are the
          results of the `inner_kernel`, `new_step_size` is a `Tensor` or a nested
          collection of `Tensor`s with the same structure as returned by the
          `step_size_getter_fn`, and `new_kernel_results` are a copy of
          `kernel_results` with the step size(s) set.
        step_size_getter_fn: A callable with the signature `(kernel_results) ->
          step_size` where `kernel_results` are the results of the `inner_kernel`,
          and `step_size` is a floating point `Tensor` or a nested collection of
          such `Tensor`s.
        log_accept_prob_getter_fn: A callable with the signature `(kernel_results)
          -> log_accept_prob` where `kernel_results` are the results of the
          `inner_kernel`, and `log_accept_prob` is a floating point `Tensor`.
          `log_accept_prob` can either be a scalar, or have shape [num_chains]. If
          it's the latter, `step_size` should also have the same leading
          dimension.
        reduce_fn: A callable with signature `(input_tensor, axis, keepdims) ->
          tensor` that returns a log-reduction of `log_accept_prob`, typically
          some sort of mean. By default, this performs an arithmetic mean.
        experimental_reduce_chain_axis_names: A `str` or list of `str`s indicating
          the named axes that should additionally reduced during the log-reduction
          of `log_accept_prob`.
        validate_args: Python `bool`. When `True` kernel parameters are checked
          for validity. When `False` invalid inputs may silently render incorrect
          outputs.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'dual_averaging_step_size_adaptation').
      """
      inner_kernel = mcmc_util.enable_store_parameters_in_results(
          inner_kernel
      )
  
      with tf.name_scope(
          mcmc_util.make_name(
              name, "dual_averaging_step_size_adaptation", "__init__"
          )
      ) as name:
          dtype = dtype_util.common_dtype(
              [
                  target_accept_prob,
                  exploration_shrinkage,
                  shrinkage_target,
                  decay_rate,
              ],
              dtype_hint=tf.float32,
          )
          target_accept_prob = tf.convert_to_tensor(
              target_accept_prob, dtype=dtype, name="target_accept_prob"
          )
          exploration_shrinkage = tf.convert_to_tensor(
              exploration_shrinkage,
              dtype=dtype,
              name="exploration_shrinkage",
          )
          step_count_smoothing = tf.cast(
              # Cast to dtype, since we asked the user to provide Int32, but we
              # want to convert to tensor here (and make it `dtype`).
              # I.e., convert_to_tensor will fail here if the user did what we
              # asked.
              step_count_smoothing,
              dtype=dtype,
              name="step_count_smoothing",
          )
          decay_rate = tf.convert_to_tensor(
              decay_rate, dtype=dtype, name="decay_rate"
          )
          num_adaptation_steps = tf.convert_to_tensor(
              num_adaptation_steps,
              dtype=tf.int32,
              name="num_adaptation_steps",
          )
          target_accept_prob = _maybe_validate_target_accept_prob(
              target_accept_prob, validate_args
          )
  
          if shrinkage_target is not None:
  
              def _convert(x):
                  return tf.convert_to_tensor(
                      x, dtype=dtype, name="shrinkage_target"
                  )
  
              shrinkage_target = tf.nest.map_structure(
                  _convert, shrinkage_target
              )
  
      self._parameters = dict(
          inner_kernel=inner_kernel,
          num_adaptation_steps=num_adaptation_steps,
          target_accept_prob=target_accept_prob,
          exploration_shrinkage=exploration_shrinkage,
          shrinkage_target=shrinkage_target,
          step_count_smoothing=step_count_smoothing,
          decay_rate=decay_rate,
          step_size_setter_fn=step_size_setter_fn,
          step_size_getter_fn=step_size_getter_fn,
          log_accept_prob_getter_fn=log_accept_prob_getter_fn,
          reduce_fn=reduce_fn,
          experimental_reduce_chain_axis_names=(
              experimental_reduce_chain_axis_names
          ),
          name=name,
      )
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow_probability.python.mcmc.simple_step_size_adaptation import get_differing_dims
  from tensorflow_probability.python.mcmc.simple_step_size_adaptation import hmc_like_log_accept_prob_getter_fn
  from tensorflow_probability.python.mcmc.simple_step_size_adaptation import hmc_like_step_size_getter_fn
  from tensorflow_probability.python.mcmc.simple_step_size_adaptation import hmc_like_step_size_setter_fn
label_desc: Use image node_id 1 to create a new DualAveragingStepSizeAdaptation object from inherited base classes: kernel_base with example: obj = DualAveragingStepSizeAdaptation(inner_kernel, num_adaptation_steps, target_accept_prob, exploration_shrinkage, shrinkage_target, step_count_smoothing, decay_rate, step_size_setter_fn, step_size_getter_fn, log_accept_prob_getter_fn, reduce_fn, experimental_reduce_chain_axis_names, validate_args, name)