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: Constructs the LeapfrogIntegrator.
  
  Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.
  
  Args:
    target_fn: Python callable which takes an argument like `*state_parts` and
      returns its (possibly unnormalized) log-density under the target
      distribution.
    step_sizes: Python `list` of `Tensor`s representing the step size for the
      leapfrog integrator. Must broadcast with the shape of
      `current_state_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    num_steps: `int` `Tensor` representing  number of steps to run
      the leapfrog integration. Total progress is roughly proportional to
      `step_size * num_steps`.
docstr_tok: ["Constructs","the","LeapfrogIntegrator",".","Assumes","a","simple","quadratic","kinetic","energy","function",":","`","0.5","||momentum||","*","*","2","`",".","Args",":","target_fn",":","Python","callable","which","takes","an","argument","like","`","*","state_parts","`","and","returns","its","(","possibly","unnormalized",")","log-density","under","the","target","distribution",".","step_sizes",":","Python","`","list","`","of","`","Tensor","`","s","representing","the","step","size","for","the","leapfrog","integrator",".","Must","broadcast","with","the","shape","of","`","current_state_parts","`",".","Larger","step","sizes","lead","to","faster","progress",",","but","too-large","step","sizes","make","rejection","exponentially","more","likely",".","When","possible",",","it","'s","often","helpful","to","match","per-variable","step","sizes","to","the","standard","deviations","of","the","target","distribution","in","each","variable",".","num_steps",":","`","int","`","`","Tensor","`","representing","number","of","steps","to","run","the","leapfrog","integration",".","Total","progress","is","roughly","proportional","to","`","step_size","*","num_steps","`","."]
code: def __init__(self, target_fn, step_sizes, num_steps):
      """Constructs the LeapfrogIntegrator.
  
      Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.
  
      Args:
        target_fn: Python callable which takes an argument like `*state_parts` and
          returns its (possibly unnormalized) log-density under the target
          distribution.
        step_sizes: Python `list` of `Tensor`s representing the step size for the
          leapfrog integrator. Must broadcast with the shape of
          `current_state_parts`.  Larger step sizes lead to faster progress, but
          too-large step sizes make rejection exponentially more likely. When
          possible, it's often helpful to match per-variable step sizes to the
          standard deviations of the target distribution in each variable.
        num_steps: `int` `Tensor` representing  number of steps to run
          the leapfrog integration. Total progress is roughly proportional to
          `step_size * num_steps`.
      """
      # Note on per-variable step sizes:
      #
      # Using per-variable step sizes is equivalent to using the same step
      # size for all variables and adding a diagonal mass matrix in the
      # kinetic energy term of the Hamiltonian being integrated. This is
      # hinted at by Neal (2011) but not derived in detail there.
      #
      # Let x and v be position and momentum variables respectively.
      # Let g(x) be the gradient of `target_fn(x)`.
      # Let S be a diagonal matrix of per-variable step sizes.
      # Let the Hamiltonian H(x, v) = -target_fn(x) + 0.5 * ||v||**2.
      #
      # Using per-variable step sizes gives the updates:
      #
      #   v' = v0 + 0.5 * S @ g(x0)
      #   x1 = x0 + S @ v'
      #   v1 = v' + 0.5 * S @ g(x1)
      #
      # Let,
      #
      #   u = inv(S) @ v
      #
      # for "u'", "u0", and "u1". Multiplying v by inv(S) in the updates above
      # gives the transformed dynamics:
      #
      #   u' = inv(S) @ v'
      #      = inv(S) @ v0 + 0.5 * g(x)
      #      = u0 + 0.5 * g(x)
      #
      #   x1 = x0 + S @ v'
      #      = x0 + S @ S @ u'
      #
      #   u1 = inv(S) @ v1
      #      = inv(S) @ v' + 0.5 * g(x1)
      #      = u' + 0.5 * g(x1)
      #
      # These are exactly the leapfrog updates for the Hamiltonian
      #
      #   H'(x, u) = -target_fn(x) + 0.5 * (S @ u).T @ (S @ u)
      #            = -target_fn(x) + 0.5 * ||v||**2
      #            = H(x, v).
      #
      # To summarize:
      #
      # * Using per-variable step sizes implicitly simulates the dynamics
      #   of the Hamiltonian H' (which are energy-conserving in H'). We
      #   keep track of v instead of u, but the underlying dynamics are
      #   the same if we transform back.
      # * The value of the Hamiltonian H'(x, u) is the same as the value
      #   of the original Hamiltonian H(x, v) after we transform back from
      #   u to v.
      # * Sampling v ~ N(0, I) is equivalent to sampling u ~ N(0, S**-2).
      #
      # So using per-variable step sizes in HMC will give results that are
      # exactly identical to explicitly using a diagonal mass matrix.
      self._target_fn = target_fn
      self._step_sizes = step_sizes
      self._num_steps = num_steps
  
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 1 to create a new SimpleLeapfrogIntegrator object from inherited base classes: LeapfrogIntegrator with example: obj = SimpleLeapfrogIntegrator(target_fn, step_sizes, num_steps)