class_docstr: Simple leapfrog integrator.
  
  Calling this functor is conceptually equivalent to:
  
  ```none
  def leapfrog(x, v, eps, L, f, M):
    g = lambda x: gradient(f, x)
    v[0] = v + eps/2 g(x)
    for l = 1...L:
      x[l] = x[l-1] + eps * inv(M) @ v[l-1]
      v[l] = v[l-1] + eps * g(x[l])
    v = v[L] - eps/2 * g(x[L])
    return x[L], v
  ```
  
  where `M = eye(dims(x))`.
  (In the future we may support arbitrary covariance `M`.)
  
  #### Examples:
  
  ```python
  import matplotlib.pyplot as plt
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator as leapfrog_impl
  dims = 10
  dtype = tf.float32
  
  target_fn = tfp.distributions.MultivariateNormalDiag(
      loc=tf.zeros(dims, dtype)).log_prob
  
  integrator = leapfrog_impl.SimpleLeapfrogIntegrator(
      target_fn,
      step_sizes=[0.1],
      num_steps=3)
  
  momentum = [tf.random.normal([dims], dtype=dtype)]
  position = [tf.random.normal([dims], dtype=dtype)]
  target = None
  target_grad_parts = None
  
  num_iter = int(1e3)
  positions = tf.zeros([num_iter, dims], dtype)
  for i in range(num_iter):
    [momentum, position, target, target_grad_parts] = integrator(
        momentum, position, target, target_grad_parts)
    positions = tf.tensor_scatter_nd_update(positions, [[i]], position)
  
  plt.plot(positions[:, 0]);  # Sinusoidal.
  ```
class_docstr_tok: ["Simple","leapfrog","integrator",".","Calling","this","functor","is","conceptually","equivalent","to",":","``","`","none","def","leapfrog","(","x",",","v",",","eps",",","L",",","f",",","M",")",":","g","=","lambda","x",":","gradient","(","f",",","x",")","v","[","0","]","=","v","+","eps\/2","g","(","x",")","for","l","=","1","...","L",":","x","[","l","]","=","x","[","l-1","]","+","eps","*","inv","(","M",")","@","v","[","l-1","]","v","[","l","]","=","v","[","l-1","]","+","eps","*","g","(","x","[","l","]",")","v","=","v","[","L","]","-","eps\/2","*","g","(","x","[","L","]",")","return","x","[","L","]",",","v","``","`","where","`","M","=","eye","(","dims","(","x",")",")","`",".","(","In","the","future","we","may","support","arbitrary","covariance","`","M","`",".",")","#","#","#","#","Examples",":","``","`","python","import","matplotlib.pyplot","as","plt","import","tensorflow.compat.v2","as","tf","import","tensorflow_probability","as","tfp","from","tensorflow_probability.python.mcmc.internal","import","leapfrog_integrator","as","leapfrog_impl","dims","=","10","dtype","=","tf.float32","target_fn","=","tfp.distributions.MultivariateNormalDiag","(","loc=tf.zeros","(","dims",",","dtype",")",")",".log_prob","integrator","=","leapfrog_impl.SimpleLeapfrogIntegrator","(","target_fn",",","step_sizes=","[","0.1","]",",","num_steps=3",")","momentum","=","[","tf.random.normal","(","[","dims","]",",","dtype=dtype",")","]","position","=","[","tf.random.normal","(","[","dims","]",",","dtype=dtype",")","]","target","=","None","target_grad_parts","=","None","num_iter","=","int","(","1e3",")","positions","=","tf.zeros","(","[","num_iter",",","dims","]",",","dtype",")","for","i","in","range","(","num_iter",")",":","[","momentum",",","position",",","target",",","target_grad_parts","]","=","integrator","(","momentum",",","position",",","target",",","target_grad_parts",")","positions","=","tf.tensor_scatter_nd_update","(","positions",",","[","[","i","]","]",",","position",")","plt.plot","(","positions","[",":",",","0","]",")",";","#","Sinusoidal",".","``","`"]
docstr: Applies `num_steps` of the leapfrog integrator.
  
  Args:
    momentum_parts: Python `list` of `Tensor`s representing momentum for each
      state part.
    state_parts: Python `list` of `Tensor`s which collectively representing
      the state.
    target: Batch of scalar `Tensor` representing the target (i.e.,
      unnormalized log prob) evaluated at `state_parts`.
    target_grad_parts: Python `list` of `Tensor`s representing the gradient of
      `target` with respect to each of `state_parts`.
    kinetic_energy_fn: Python callable that can evaluate the kinetic energy
      of the given momentum. This is typically the negative log probability of
      the distribution over the momentum.
    name: Python `str` used to group ops created by this function.
  
  Returns:
    next_momentum_parts: Python `list` of `Tensor`s representing new momentum.
    next_state_parts: Python `list` of `Tensor`s which collectively
      representing the new state.
    next_target: Batch of scalar `Tensor` representing the target (i.e.,
      unnormalized log prob) evaluated at `next_state_parts`.
    next_target_grad_parts: Python `list` of `Tensor`s representing the
      gradient of `next_target` with respect to each of `next_state_parts`.
docstr_tok: ["Applies","`","num_steps","`","of","the","leapfrog","integrator",".","Args",":","momentum_parts",":","Python","`","list","`","of","`","Tensor","`","s","representing","momentum","for","each","state","part",".","state_parts",":","Python","`","list","`","of","`","Tensor","`","s","which","collectively","representing","the","state",".","target",":","Batch","of","scalar","`","Tensor","`","representing","the","target","(","i.e.",",","unnormalized","log","prob",")","evaluated","at","`","state_parts","`",".","target_grad_parts",":","Python","`","list","`","of","`","Tensor","`","s","representing","the","gradient","of","`","target","`","with","respect","to","each","of","`","state_parts","`",".","kinetic_energy_fn",":","Python","callable","that","can","evaluate","the","kinetic","energy","of","the","given","momentum",".","This","is","typically","the","negative","log","probability","of","the","distribution","over","the","momentum",".","name",":","Python","`","str","`","used","to","group","ops","created","by","this","function",".","Returns",":","next_momentum_parts",":","Python","`","list","`","of","`","Tensor","`","s","representing","new","momentum",".","next_state_parts",":","Python","`","list","`","of","`","Tensor","`","s","which","collectively","representing","the","new","state",".","next_target",":","Batch","of","scalar","`","Tensor","`","representing","the","target","(","i.e.",",","unnormalized","log","prob",")","evaluated","at","`","next_state_parts","`",".","next_target_grad_parts",":","Python","`","list","`","of","`","Tensor","`","s","representing","the","gradient","of","`","next_target","`","with","respect","to","each","of","`","next_state_parts","`","."]
code: def __call__(
      self,
      momentum_parts,
      state_parts,
      target=None,
      target_grad_parts=None,
      kinetic_energy_fn=None,
      name=None,
  ):
      """Applies `num_steps` of the leapfrog integrator.
  
      Args:
        momentum_parts: Python `list` of `Tensor`s representing momentum for each
          state part.
        state_parts: Python `list` of `Tensor`s which collectively representing
          the state.
        target: Batch of scalar `Tensor` representing the target (i.e.,
          unnormalized log prob) evaluated at `state_parts`.
        target_grad_parts: Python `list` of `Tensor`s representing the gradient of
          `target` with respect to each of `state_parts`.
        kinetic_energy_fn: Python callable that can evaluate the kinetic energy
          of the given momentum. This is typically the negative log probability of
          the distribution over the momentum.
        name: Python `str` used to group ops created by this function.
  
      Returns:
        next_momentum_parts: Python `list` of `Tensor`s representing new momentum.
        next_state_parts: Python `list` of `Tensor`s which collectively
          representing the new state.
        next_target: Batch of scalar `Tensor` representing the target (i.e.,
          unnormalized log prob) evaluated at `next_state_parts`.
        next_target_grad_parts: Python `list` of `Tensor`s representing the
          gradient of `next_target` with respect to each of `next_state_parts`.
      """
      with tf.name_scope(name or "leapfrog_integrate"):
          [
              momentum_parts,
              state_parts,
              target,
              target_grad_parts,
          ] = process_args(
              self.target_fn,
              momentum_parts,
              state_parts,
              target,
              target_grad_parts,
          )
  
          if kinetic_energy_fn is None:
              # Avoid adding ops and taking grads, when the implied kinetic energy
              # is just 0.5 * ||x||^2, so the gradient is x
              get_velocity_parts = lambda x: x
          else:
  
              def get_velocity_parts(half_next_momentum_parts):
                  _, velocity_parts = mcmc_util.maybe_call_fn_and_grads(
                      kinetic_energy_fn, half_next_momentum_parts
                  )
                  return velocity_parts
  
          # See Algorithm 1 of "Faster Hamiltonian Monte Carlo by Learning Leapfrog
          # Scale", https://arxiv.org/abs/1810.04449.
  
          half_next_momentum_parts = [
              v + _multiply(0.5 * eps, g, dtype=v.dtype)
              for v, eps, g in zip(
                  momentum_parts, self.step_sizes, target_grad_parts
              )
          ]
  
          [
              _,
              next_half_next_momentum_parts,
              next_state_parts,
              next_target,
              next_target_grad_parts,
          ] = tf.while_loop(
              cond=lambda i, *_: i < self.num_steps,
              body=lambda i, *args: [i + 1]
              + list(
                  _one_step(  # pylint: disable=no-value-for-parameter,g-long-lambda
                      self.target_fn,
                      self.step_sizes,
                      get_velocity_parts,
                      *args,
                  )
              ),
              loop_vars=[
                  tf.zeros_like(self.num_steps, name="iter"),
                  half_next_momentum_parts,
                  state_parts,
                  target,
                  target_grad_parts,
              ],
          )
  
          next_momentum_parts = [
              v
              - _multiply(
                  0.5 * eps, g, dtype=v.dtype
              )  # pylint: disable=g-complex-comprehension
              for v, eps, g in zip(
                  next_half_next_momentum_parts,
                  self.step_sizes,
                  next_target_grad_parts,
              )
          ]
  
          return (
              next_momentum_parts,
              next_state_parts,
              next_target,
              next_target_grad_parts,
          )
  
file_path: probability/tensorflow_probability/python/mcmc/internal/leapfrog_integrator.py
imports: import abc
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 5 for calling the SimpleLeapfrogIntegrator obj's underlying member method code with example usage: obj.__call__(momentum_parts, state_parts, target, target_grad_parts, kinetic_energy_fn, name) and returns: next_momentum_parts,  next_state_parts,  next_target,  next_target_grad_parts, velocity_parts