code: def _leapfrog(
      value_and_gradients_fn,
      current,
      step_size,
      direction,
      unrolled_leapfrog_steps,
  ):
      """_leapfrog_base, with input rank padding by force of will."""
      # The purpose of this padding is to simulate broadcasting rank expansion
      # despite a batch dimension.
  
      # The problem is this: Consider, say, adding two Tensors of shape [2, 3] and
      # [3].  Addition broadcasts, which automatically expands the rank of the
      # second one by padding with size-1 dimensions, to [1, 3], and matches the [2,
      # 3] shape by tiling.  But, what happens if both of these Tensors gain a batch
      # dimension of size 4?  Naively, the shapes become [4, 2, 3] and [4, 3].  The
      # batch dimension of the second Tensor now aligns with the top non-batch
      # dimension of the first, and broadcasting fails.
  
      # This happens in the leapfrog integrator for matching step sizes with the
      # state Tensors being stepped.  It also happens as a knock-on for the
      # direction of integration, which multiplies the step size.
      step_size = [
          _expand_dims_under_batch_dim(step, tf.rank(state))
          for step, state in zip(step_size, current.state)
      ]
      direction = [
          _expand_dims_under_batch_dim(direction, tf.rank(step))
          for step in step_size
      ]
      return _leapfrog_base(
          value_and_gradients_fn,
          current,
          step_size,
          direction,
          unrolled_leapfrog_steps,
      )
  
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 7 for calling a global function with example usage: _leapfrog(value_and_gradients_fn, current, step_size, direction, unrolled_leapfrog_steps) and returns: _leapfrog_base