class_docstr: Adapts the inner kernel's `step_size` based on `log_accept_prob`.
  
  The simple policy multiplicatively increases or decreases the `step_size` of
  the inner kernel based on the value of `log_accept_prob`. It is based on
  [equation 19 of Andrieu and Thoms (2008)][1]. Given enough steps and small
  enough `adaptation_rate` the median of the distribution of the acceptance
  probability will converge to the `target_accept_prob`. 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][2] 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, Y]` (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` [3]. 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.SimpleStepSizeAdaptation(
      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]: Andrieu, Christophe, Thoms, Johannes. A tutorial on adaptive MCMC.
       _Statistics and Computing_, 2008.
       https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf
  
  [2]:
  http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745
  
  [3]: Hoffman, M., Radul, A., & Sountsov, P. An Adaptive MCMC Scheme
       for Setting Trajectory Lengths in Hamiltonian Monte Carlo, 2020. In
       preparation.
class_docstr_tok: ["Adapts","the","inner","kernel","'s","`","step_size","`","based","on","`","log_accept_prob","`",".","The","simple","policy","multiplicatively","increases","or","decreases","the","`","step_size","`","of","the","inner","kernel","based","on","the","value","of","`","log_accept_prob","`",".","It","is","based","on","[","equation","19","of","Andrieu","and","Thoms","(","2008",")","]","[","1","]",".","Given","enough","steps","and","small","enough","`","adaptation_rate","`","the","median","of","the","distribution","of","the","acceptance","probability","will","converge","to","the","`","target_accept_prob","`",".","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","]","[","2","]","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",",","Y","]","`","(","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","do","n'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","`","[","3","]",".","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.SimpleStepSizeAdaptation","(","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","]",":","Andrieu",",","Christophe",",","Thoms",",","Johannes",".","A","tutorial","on","adaptive","MCMC",".","_Statistics","and","Computing_",",","2008.","https",":","\/\/people.eecs.berkeley.edu\/~jordan\/sail\/readings\/andrieu-thoms.pdf","[","2","]",":","http",":","\/\/andrewgelman.com\/2017\/12\/15\/burn-vs-warm-iterative-simulation-algorithms\/","#","comment-627745","[","3","]",":","Hoffman",",","M.",",","Radul",",","A.",",","&","Sountsov",",","P.","An","Adaptive","MCMC","Scheme","for","Setting","Trajectory","Lengths","in","Hamiltonian","Monte","Carlo",",","2020",".","In","preparation","."]
code: def name(self):
      return self._parameters["name"]
  
imports: import collections
  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.internal import tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.math.generic import reduce_logmeanexp
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 3 for calling the SimpleStepSizeAdaptation obj's underlying member method code with example usage: obj.name() and returns: self