def leapfrog_step(leapfrog_step_state: LeapFrogStepState,<EOL>step_size: FloatTensor, target_log_prob_fn: PotentialFn,<EOL>kinetic_energy_fn: PotentialFn<EOL>) -> Tuple[LeapFrogStepState, LeapFrogStepExtras]:
state = leapfrog_step_state.state<EOL>state_grads = leapfrog_step_state.state_grads<EOL>momentum = leapfrog_step_state.momentum<EOL>step_size = maybe_broadcast_structure(step_size, state)<EOL>state = tf.nest.map_structure(tf.convert_to_tensor, state)<EOL>momentum = tf.nest.map_structure(tf.convert_to_tensor, momentum)<EOL>state = tf.nest.map_structure(tf.convert_to_tensor, state)<EOL>if state_grads is None:<EOL><INDENT>_, _, state_grads = call_and_grads(target_log_prob_fn, state)<EOL><DEDENT>else:<EOL><INDENT>state_grads = tf.nest.map_structure(tf.convert_to_tensor, state_grads)<EOL><DEDENT>momentum = tf.nest.map_structure(lambda m, sg, s: m + <NUM_LIT:0.5> * sg * s, momentum,<EOL>state_grads, step_size)<EOL>kinetic_energy, kinetic_energy_extra, momentum_grads = call_and_grads(<EOL>kinetic_energy_fn, momentum)<EOL>state = tf.nest.map_structure(lambda x, mg, s: x + mg * s, state,<EOL>momentum_grads, step_size)<EOL>target_log_prob, state_extra, state_grads = call_and_grads(<EOL>target_log_prob_fn, state)<EOL>momentum = tf.nest.map_structure(lambda m, sg, s: m + <NUM_LIT:0.5> * sg * s, momentum,<EOL>state_grads, step_size)<EOL>return LeapFrogStepState(state, state_grads, momentum), LeapFrogStepExtras(<EOL>target_log_prob, state_extra, kinetic_energy, kinetic_energy_extra)<EOL>
Leapfrog `TransitionOperator`.

    Args:
      leapfrog_step_state: LeapFrogStepState.
      step_size: Step size, structure broadcastable to the `target_log_prob_fn`
        state.
      target_log_prob_fn: Target log prob fn.
      kinetic_energy_fn: Kinetic energy fn.

    Returns:
      leapfrog_step_state: LeapFrogStepState.
      leapfrog_step_extras: LeapFrogStepExtras.