class_docstr: Runs one step of Hamiltonian Monte Carlo.
  
  Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) algorithm
  that takes a series of gradient-informed steps to produce a Metropolis
  proposal. This class implements one random HMC step from a given
  `current_state`. Mathematical details and derivations can be found in
  [Neal (2011)][1].
  
  The `one_step` function can update multiple chains in parallel. It assumes
  that all leftmost dimensions of `current_state` index independent chain states
  (and are therefore updated independently). The output of
  `target_log_prob_fn(*current_state)` should sum log-probabilities across all
  event dimensions. Slices along the rightmost dimensions may have different
  target distributions; for example, `current_state[0, :]` could have a
  different target distribution from `current_state[1, :]`. These semantics are
  governed by `target_log_prob_fn(*current_state)`. (The number of independent
  chains is `tf.size(target_log_prob_fn(*current_state))`.)
  
  #### Examples:
  
  ##### Simple chain with warm-up.
  
  In this example we sample from a standard univariate normal
  distribution using HMC with adaptive step size.
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  
  tf.enable_eager_execution()
  
  # Target distribution is proportional to: `exp(-x (1 + x))`.
  def unnormalized_log_prob(x):
    return -x - x**2.
  
  # Initialize the HMC transition kernel.
  num_results = int(10e3)
  num_burnin_steps = int(1e3)
  adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
      tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=unnormalized_log_prob,
          num_leapfrog_steps=3,
          step_size=1.),
      num_adaptation_steps=int(num_burnin_steps * 0.8))
  
  # Run the chain (with burn-in).
  @tf.function
  def run_chain():
    # Run the chain (with burn-in).
    samples, is_accepted = tfp.mcmc.sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=1.,
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted)
  
    sample_mean = tf.reduce_mean(samples)
    sample_stddev = tf.math.reduce_std(samples)
    is_accepted = tf.reduce_mean(tf.cast(is_accepted, dtype=tf.float32))
    return sample_mean, sample_stddev, is_accepted
  
  sample_mean, sample_stddev, is_accepted = run_chain()
  
  print('mean:{:.4f}  stddev:{:.4f}  acceptance:{:.4f}'.format(
      sample_mean.numpy(), sample_stddev.numpy(), is_accepted.numpy()))
  ```
  
  ##### Estimate parameters of a more complicated posterior.
  
  In this example, we'll use Monte-Carlo EM to find best-fit parameters. See
  [_Convergence of a stochastic approximation version of the EM algorithm_][2]
  for more details.
  
  More precisely, we use HMC to form a chain conditioned on parameter `sigma`
  and training data `{ (x[i], y[i]) : i=1...n }`. Then we use one gradient step
  of maximum-likelihood to improve the `sigma` estimate. Then repeat the process
  until convergence. (This procedure is a [Robbins--Monro algorithm](
  https://en.wikipedia.org/wiki/Stochastic_approximation).)
  
  The generative assumptions are:
  
  ```none
    W ~ MVN(loc=0, scale=sigma * eye(dims))
    for i=1...num_samples:
        X[i] ~ MVN(loc=0, scale=eye(dims))
      eps[i] ~ Normal(loc=0, scale=1)
        Y[i] = X[i].T * W + eps[i]
  ```
  
  We now implement a stochastic approximation of Expectation Maximization (SAEM)
  using `tensorflow_probability` intrinsics. [Bernard (1999)][2]
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  import numpy as np
  
  tf.enable_eager_execution()
  
  tfd = tfp.distributions
  
  def make_training_data(num_samples, dims, sigma):
    dt = np.asarray(sigma).dtype
    x = np.random.randn(dims, num_samples).astype(dt)
    w = sigma * np.random.randn(1, dims).astype(dt)
    noise = np.random.randn(num_samples).astype(dt)
    y = w.dot(x) + noise
    return y[0], x, w[0]
  
  def make_weights_prior(dims, log_sigma):
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros([dims], dtype=log_sigma.dtype),
        scale_diag=tf.math.exp(log_sigma) *
                   tf.ones([dims], dtype=log_sigma.dtype))
  
  def make_response_likelihood(w, x):
    if w.shape.ndims == 1:
      y_bar = tf.matmul(w[tf.newaxis], x)[0]
    else:
      y_bar = tf.matmul(w, x)
    return tfd.Normal(loc=y_bar, scale=tf.ones_like(y_bar))  # [n]
  
  # Setup assumptions.
  dtype = np.float32
  num_samples = 500
  dims = 10
  tf.random.set_seed(10014)
  np.random.seed(10014)
  
  weights_prior_true_scale = np.array(0.3, dtype)
  y, x, _ = make_training_data(
      num_samples, dims, weights_prior_true_scale)
  
  log_sigma = tf.Variable(0., dtype=dtype, name='log_sigma')
  
  optimizer = tf_keras.optimizers.SGD(learning_rate=0.01)
  
  @tf.function
  def mcem_iter(weights_chain_start, step_size):
    with tf.GradientTape() as tape:
      tape.watch(log_sigma)
      prior = make_weights_prior(dims, log_sigma)
  
      def unnormalized_posterior_log_prob(w):
        likelihood = make_response_likelihood(w, x)
        return (
            prior.log_prob(w) +
            tf.reduce_sum(likelihood.log_prob(y), axis=-1))  # [m]
  
      def trace_fn(_, pkr):
        return (
            pkr.inner_results.log_accept_ratio,
            pkr.inner_results.accepted_results.target_log_prob,
            pkr.inner_results.accepted_results.step_size)
  
      num_results = 2
      weights, (
          log_accept_ratio, target_log_prob, step_size) = tfp.mcmc.sample_chain(
          num_results=num_results,
          num_burnin_steps=0,
          current_state=weights_chain_start,
          kernel=tfp.mcmc.SimpleStepSizeAdaptation(
              tfp.mcmc.HamiltonianMonteCarlo(
                  target_log_prob_fn=unnormalized_posterior_log_prob,
                  num_leapfrog_steps=2,
                  step_size=step_size,
                  state_gradients_are_stopped=True,
              ),
              # Adapt for the entirety of the trajectory.
              num_adaptation_steps=2),
          trace_fn=trace_fn,
          seed=123)
  
      # We do an optimization step to propagate `log_sigma` after two HMC
      # steps to propagate `weights`.
      loss = -tf.reduce_mean(target_log_prob)
  
    avg_acceptance_ratio = tf.math.exp(
        tfp.math.reduce_logmeanexp(tf.minimum(log_accept_ratio, 0.)))
  
    optimizer.apply_gradients(
        [[tape.gradient(loss, log_sigma), log_sigma]])
  
    weights_prior_estimated_scale = tf.math.exp(log_sigma)
    return (weights_prior_estimated_scale, weights[-1], loss,
            step_size[-1], avg_acceptance_ratio)
  
  num_iters = int(40)
  
  weights_prior_estimated_scale_ = np.zeros(num_iters, dtype)
  weights_ = np.zeros([num_iters + 1, dims], dtype)
  loss_ = np.zeros([num_iters], dtype)
  weights_[0] = np.random.randn(dims).astype(dtype)
  step_size_ = 0.03
  
  for iter_ in range(num_iters):
    [
        weights_prior_estimated_scale_[iter_],
        weights_[iter_ + 1],
        loss_[iter_],
        step_size_,
        avg_acceptance_ratio_,
    ] = mcem_iter(weights_[iter_], step_size_)
    tf.compat.v1.logging.vlog(
        1, ('iter:{:>2}  loss:{: 9.3f}  scale:{:.3f}  '
            'step_size:{:.4f}  avg_acceptance_ratio:{:.4f}').format(
                iter_, loss_[iter_], weights_prior_estimated_scale_[iter_],
                step_size_, avg_acceptance_ratio_))
  
  # Should converge to ~0.22.
  import matplotlib.pyplot as plt
  plt.plot(weights_prior_estimated_scale_)
  plt.ylabel('weights_prior_estimated_scale')
  plt.xlabel('iteration')
  ```
  
  #### References
  
  [1]: Radford Neal. MCMC Using Hamiltonian Dynamics. _Handbook of Markov Chain
       Monte Carlo_, 2011. https://arxiv.org/abs/1206.1901
  
  [2]: Bernard Delyon, Marc Lavielle, Eric, Moulines. _Convergence of a
       stochastic approximation version of the EM algorithm_, Ann. Statist. 27
       (1999), no. 1, 94--128. https://projecteuclid.org/euclid.aos/1018031103
class_docstr_tok: ["Runs","one","step","of","Hamiltonian","Monte","Carlo",".","Hamiltonian","Monte","Carlo","(","HMC",")","is","a","Markov","chain","Monte","Carlo","(","MCMC",")","algorithm","that","takes","a","series","of","gradient-informed","steps","to","produce","a","Metropolis","proposal",".","This","class","implements","one","random","HMC","step","from","a","given","`","current_state","`",".","Mathematical","details","and","derivations","can","be","found","in","[","Neal","(","2011",")","]","[","1","]",".","The","`","one_step","`","function","can","update","multiple","chains","in","parallel",".","It","assumes","that","all","leftmost","dimensions","of","`","current_state","`","index","independent","chain","states","(","and","are","therefore","updated","independently",")",".","The","output","of","`","target_log_prob_fn","(","*","current_state",")","`","should","sum","log-probabilities","across","all","event","dimensions",".","Slices","along","the","rightmost","dimensions","may","have","different","target","distributions",";","for","example",",","`","current_state","[","0",",",":","]","`","could","have","a","different","target","distribution","from","`","current_state","[","1",",",":","]","`",".","These","semantics","are","governed","by","`","target_log_prob_fn","(","*","current_state",")","`",".","(","The","number","of","independent","chains","is","`","tf.size","(","target_log_prob_fn","(","*","current_state",")",")","`",".",")","#","#","#","#","Examples",":","#","#","#","#","#","Simple","chain","with","warm-up",".","In","this","example","we","sample","from","a","standard","univariate","normal","distribution","using","HMC","with","adaptive","step","size",".","``","`","python","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","tf.enable_eager_execution","(",")","#","Target","distribution","is","proportional","to",":","`","exp","(","-x","(","1","+","x",")",")","`",".","def","unnormalized_log_prob","(","x",")",":","return","-x","-","x","*","*","2",".","#","Initialize","the","HMC","transition","kernel",".","num_results","=","int","(","10e3",")","num_burnin_steps","=","int","(","1e3",")","adaptive_hmc","=","tfp.mcmc.SimpleStepSizeAdaptation","(","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=unnormalized_log_prob",",","num_leapfrog_steps=3",",","step_size=1",".",")",",","num_adaptation_steps=int","(","num_burnin_steps","*","0.8",")",")","#","Run","the","chain","(","with","burn-in",")",".","@","tf.function","def","run_chain","(",")",":","#","Run","the","chain","(","with","burn-in",")",".","samples",",","is_accepted","=","tfp.mcmc.sample_chain","(","num_results=num_results",",","num_burnin_steps=num_burnin_steps",",","current_state=1.",",","kernel=adaptive_hmc",",","trace_fn=lambda","_",",","pkr",":","pkr.inner_results.is_accepted",")","sample_mean","=","tf.reduce_mean","(","samples",")","sample_stddev","=","tf.math.reduce_std","(","samples",")","is_accepted","=","tf.reduce_mean","(","tf.cast","(","is_accepted",",","dtype=tf.float32",")",")","return","sample_mean",",","sample_stddev",",","is_accepted","sample_mean",",","sample_stddev",",","is_accepted","=","run_chain","(",")","print","(","'mean",":","{",":",".4f","}","stddev",":","{",":",".4f","}","acceptance",":","{",":",".4f","}","'.format","(","sample_mean.numpy","(",")",",","sample_stddev.numpy","(",")",",","is_accepted.numpy","(",")",")",")","``","`","#","#","#","#","#","Estimate","parameters","of","a","more","complicated","posterior",".","In","this","example",",","we","'ll","use","Monte-Carlo","EM","to","find","best-fit","parameters",".","See","[","_Convergence","of","a","stochastic","approximation","version","of","the","EM","algorithm_","]","[","2","]","for","more","details",".","More","precisely",",","we","use","HMC","to","form","a","chain","conditioned","on","parameter","`","sigma","`","and","training","data","`","{","(","x","[","i","]",",","y","[","i","]",")",":","i=1","...","n","}","`",".","Then","we","use","one","gradient","step","of","maximum-likelihood","to","improve","the","`","sigma","`","estimate",".","Then","repeat","the","process","until","convergence",".","(","This","procedure","is","a","[","Robbins","--","Monro","algorithm","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Stochastic_approximation",")",".",")","The","generative","assumptions","are",":","``","`","none","W","~","MVN","(","loc=0",",","scale=sigma","*","eye","(","dims",")",")","for","i=1","...","num_samples",":","X","[","i","]","~","MVN","(","loc=0",",","scale=eye","(","dims",")",")","eps","[","i","]","~","Normal","(","loc=0",",","scale=1",")","Y","[","i","]","=","X","[","i","]",".T","*","W","+","eps","[","i","]","``","`","We","now","implement","a","stochastic","approximation","of","Expectation","Maximization","(","SAEM",")","using","`","tensorflow_probability","`","intrinsics",".","[","Bernard","(","1999",")","]","[","2","]","``","`","python","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","import","numpy","as","np","tf.enable_eager_execution","(",")","tfd","=","tfp.distributions","def","make_training_data","(","num_samples",",","dims",",","sigma",")",":","dt","=","np.asarray","(","sigma",")",".dtype","x","=","np.random.randn","(","dims",",","num_samples",")",".astype","(","dt",")","w","=","sigma","*","np.random.randn","(","1",",","dims",")",".astype","(","dt",")","noise","=","np.random.randn","(","num_samples",")",".astype","(","dt",")","y","=","w.dot","(","x",")","+","noise","return","y","[","0","]",",","x",",","w","[","0","]","def","make_weights_prior","(","dims",",","log_sigma",")",":","return","tfd.MultivariateNormalDiag","(","loc=tf.zeros","(","[","dims","]",",","dtype=log_sigma.dtype",")",",","scale_diag=tf.math.exp","(","log_sigma",")","*","tf.ones","(","[","dims","]",",","dtype=log_sigma.dtype",")",")","def","make_response_likelihood","(","w",",","x",")",":","if","w.shape.ndims","==","1",":","y_bar","=","tf.matmul","(","w","[","tf.newaxis","]",",","x",")","[","0","]","else",":","y_bar","=","tf.matmul","(","w",",","x",")","return","tfd.Normal","(","loc=y_bar",",","scale=tf.ones_like","(","y_bar",")",")","#","[","n","]","#","Setup","assumptions",".","dtype","=","np.float32","num_samples","=","500","dims","=","10","tf.random.set_seed","(","10014",")","np.random.seed","(","10014",")","weights_prior_true_scale","=","np.array","(","0.3",",","dtype",")","y",",","x",",","_","=","make_training_data","(","num_samples",",","dims",",","weights_prior_true_scale",")","log_sigma","=","tf.Variable","(","0.",",","dtype=dtype",",","name='log_sigma","'",")","optimizer","=","tf_keras.optimizers.SGD","(","learning_rate=0.01",")","@","tf.function","def","mcem_iter","(","weights_chain_start",",","step_size",")",":","with","tf.GradientTape","(",")","as","tape",":","tape.watch","(","log_sigma",")","prior","=","make_weights_prior","(","dims",",","log_sigma",")","def","unnormalized_posterior_log_prob","(","w",")",":","likelihood","=","make_response_likelihood","(","w",",","x",")","return","(","prior.log_prob","(","w",")","+","tf.reduce_sum","(","likelihood.log_prob","(","y",")",",","axis=-1",")",")","#","[","m","]","def","trace_fn","(","_",",","pkr",")",":","return","(","pkr.inner_results.log_accept_ratio",",","pkr.inner_results.accepted_results.target_log_prob",",","pkr.inner_results.accepted_results.step_size",")","num_results","=","2","weights",",","(","log_accept_ratio",",","target_log_prob",",","step_size",")","=","tfp.mcmc.sample_chain","(","num_results=num_results",",","num_burnin_steps=0",",","current_state=weights_chain_start",",","kernel=tfp.mcmc.SimpleStepSizeAdaptation","(","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=unnormalized_posterior_log_prob",",","num_leapfrog_steps=2",",","step_size=step_size",",","state_gradients_are_stopped=True",",",")",",","#","Adapt","for","the","entirety","of","the","trajectory",".","num_adaptation_steps=2",")",",","trace_fn=trace_fn",",","seed=123",")","#","We","do","an","optimization","step","to","propagate","`","log_sigma","`","after","two","HMC","#","steps","to","propagate","`","weights","`",".","loss","=","-tf.reduce_mean","(","target_log_prob",")","avg_acceptance_ratio","=","tf.math.exp","(","tfp.math.reduce_logmeanexp","(","tf.minimum","(","log_accept_ratio",",","0",".",")",")",")","optimizer.apply_gradients","(","[","[","tape.gradient","(","loss",",","log_sigma",")",",","log_sigma","]","]",")","weights_prior_estimated_scale","=","tf.math.exp","(","log_sigma",")","return","(","weights_prior_estimated_scale",",","weights","[","-1","]",",","loss",",","step_size","[","-1","]",",","avg_acceptance_ratio",")","num_iters","=","int","(","40",")","weights_prior_estimated_scale_","=","np.zeros","(","num_iters",",","dtype",")","weights_","=","np.zeros","(","[","num_iters","+","1",",","dims","]",",","dtype",")","loss_","=","np.zeros","(","[","num_iters","]",",","dtype",")","weights_","[","0","]","=","np.random.randn","(","dims",")",".astype","(","dtype",")","step_size_","=","0.03","for","iter_","in","range","(","num_iters",")",":","[","weights_prior_estimated_scale_","[","iter_","]",",","weights_","[","iter_","+","1","]",",","loss_","[","iter_","]",",","step_size_",",","avg_acceptance_ratio_",",","]","=","mcem_iter","(","weights_","[","iter_","]",",","step_size_",")","tf.compat.v1.logging.vlog","(","1",",","(","'iter",":","{",":",">","2","}","loss",":","{",":","9.3f","}","scale",":","{",":",".3f","}","'","'step_size",":","{",":",".4f","}","avg_acceptance_ratio",":","{",":",".4f","}","'",")",".format","(","iter_",",","loss_","[","iter_","]",",","weights_prior_estimated_scale_","[","iter_","]",",","step_size_",",","avg_acceptance_ratio_",")",")","#","Should","converge","to","~0.22",".","import","matplotlib.pyplot","as","plt","plt.plot","(","weights_prior_estimated_scale_",")","plt.ylabel","(","'weights_prior_estimated_scale","'",")","plt.xlabel","(","'iteration","'",")","``","`","#","#","#","#","References","[","1","]",":","Radford","Neal",".","MCMC","Using","Hamiltonian","Dynamics",".","_Handbook","of","Markov","Chain","Monte","Carlo_",",","2011.","https",":","\/\/arxiv.org\/abs\/1206.1901","[","2","]",":","Bernard","Delyon",",","Marc","Lavielle",",","Eric",",","Moulines",".","_Convergence","of","a","stochastic","approximation","version","of","the","EM","algorithm_",",","Ann",".","Statist",".","27","(","1999",")",",","no",".","1",",","94","--","128.","https",":","\/\/projecteuclid.org\/euclid.aos\/1018031103"]
docstr: Returns the num_leapfrog_steps parameter.
  
  If `store_parameters_in_results` argument to the initializer was set to
  `True`, this only returns the value of the `num_leapfrog_steps` placed in
  the kernel results by the `bootstrap_results` method. The actual
  `num_leapfrog_steps` in that situation is governed by the
  `previous_kernel_results` argument to `one_step` method.
  
  Returns:
    num_leapfrog_steps: An integer `Tensor`.
docstr_tok: ["Returns","the","num_leapfrog_steps","parameter",".","If","`","store_parameters_in_results","`","argument","to","the","initializer","was","set","to","`","True","`",",","this","only","returns","the","value","of","the","`","num_leapfrog_steps","`","placed","in","the","kernel","results","by","the","`","bootstrap_results","`","method",".","The","actual","`","num_leapfrog_steps","`","in","that","situation","is","governed","by","the","`","previous_kernel_results","`","argument","to","`","one_step","`","method",".","Returns",":","num_leapfrog_steps",":","An","integer","`","Tensor","`","."]
code: def num_leapfrog_steps(self):
      """Returns the num_leapfrog_steps parameter.
  
      If `store_parameters_in_results` argument to the initializer was set to
      `True`, this only returns the value of the `num_leapfrog_steps` placed in
      the kernel results by the `bootstrap_results` method. The actual
      `num_leapfrog_steps` in that situation is governed by the
      `previous_kernel_results` argument to `one_step` method.
  
      Returns:
        num_leapfrog_steps: An integer `Tensor`.
      """
      return self._impl.inner_kernel.num_leapfrog_steps
  
imports: import collections
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 4 for calling the HamiltonianMonteCarlo obj's underlying member method code with example usage: obj.num_leapfrog_steps() and returns: self