code: def _one_step(
      target_fn,
      step_sizes,
      get_velocity_parts,
      half_next_momentum_parts,
      state_parts,
      target,
      target_grad_parts,
  ):
      """Body of integrator while loop."""
      with tf.name_scope("leapfrog_integrate_one_step"):
          velocity_parts = get_velocity_parts(half_next_momentum_parts)
          next_state_parts = []
          for state_part, eps, velocity_part in zip(
              state_parts, step_sizes, velocity_parts
          ):
              next_state_parts.append(
                  state_part
                  + _multiply(
                      eps, velocity_part, dtype=state_part.dtype
                  )
              )
          [
              next_target,
              next_target_grad_parts,
          ] = mcmc_util.maybe_call_fn_and_grads(
              target_fn, next_state_parts
          )
          if any(g is None for g in next_target_grad_parts):
              raise ValueError(
                  "Encountered `None` gradient.\n"
                  "  state_parts: {}\n"
                  "  next_state_parts: {}\n"
                  "  next_target_grad_parts: {}".format(
                      state_parts,
                      next_state_parts,
                      next_target_grad_parts,
                  )
              )
  
          tensorshape_util.set_shape(next_target, target.shape)
          for ng, g in zip(next_target_grad_parts, target_grad_parts):
              tensorshape_util.set_shape(ng, g.shape)
  
          next_half_next_momentum_parts = [
              v
              + _multiply(
                  eps, g, dtype=v.dtype
              )  # pylint: disable=g-complex-comprehension
              for v, eps, g in zip(
                  half_next_momentum_parts,
                  step_sizes,
                  next_target_grad_parts,
              )
          ]
  
          return [
              next_half_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 1 for calling a global function with example usage: _one_step(target_fn, step_sizes, get_velocity_parts, half_next_momentum_parts, state_parts, target, target_grad_parts) and returns: list