code: def _leapfrog_base(
      value_and_gradients_fn,
      current,
      step_size,
      direction,
      unrolled_leapfrog_steps,
  ):
      """Runs `unrolled_leapfrog_steps` steps of leapfrog integration."""
      with tf.name_scope("leapfrog"):
          step_size = [d * s for d, s in zip(direction, step_size)]
          for _ in range(unrolled_leapfrog_steps):
              mid_momentum = [
                  m + 0.5 * step * g
                  for m, step, g in zip(
                      current.momentum,
                      step_size,
                      current.grads_target_log_prob,
                  )
              ]
              next_state = [
                  s + step * m
                  for s, step, m in zip(
                      current.state, step_size, mid_momentum
                  )
              ]
              with tf.name_scope("gradients"):
                  [
                      next_target_log_prob,
                      next_grads_target_log_prob,
                  ] = value_and_gradients_fn(*next_state)
              next_momentum = [
                  m + 0.5 * step * g
                  for m, step, g in zip(
                      mid_momentum,
                      step_size,
                      next_grads_target_log_prob,
                  )
              ]
              current = Point(
                  next_state,
                  next_target_log_prob,
                  next_grads_target_log_prob,
                  next_momentum,
              )
          return current
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python import random
  from tensorflow_probability.python.experimental import auto_batching
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.util.seed_stream import SeedStream
label_desc: Use image node_id 6 for calling a global function with example usage: _leapfrog_base(value_and_gradients_fn, current, step_size, direction, unrolled_leapfrog_steps) and returns: current