code: def has_not_u_turn_at_all_index(
      read_indexes,
      direction,
      momentum_state_memory,
      momentum_right,
      state_right,
      no_u_turns_within_tree,
      log_prob_rank,
      shard_axis_names=None,
  ):
      """Check u turn for early stopping."""
  
      def _get_left_state_and_check_u_turn(
          left_current_index, no_u_turns_last
      ):
          """Check U turn on a single index."""
          momentum_left = [
              tf.gather(x, left_current_index, axis=0)
              for x in momentum_state_memory.momentum_swap
          ]
          state_left = [
              tf.gather(x, left_current_index, axis=0)
              for x in momentum_state_memory.state_swap
          ]
          # Note that in generalized u turn, state_diff is actually the cumulated sum
          # of the momentum.
          state_diff = [
              s1 - s2 for s1, s2 in zip(state_right, state_left)
          ]
          if not GENERALIZED_UTURN:
              state_diff = [
                  tf.where(d, m, -m)
                  for d, m in zip(direction, state_diff)
              ]
  
          no_u_turns_current = has_not_u_turn(
              state_diff,
              momentum_left,
              momentum_right,
              log_prob_rank,
              shard_axis_names=shard_axis_names,
          )
          return (
              left_current_index + 1,
              no_u_turns_current & no_u_turns_last,
          )
  
      _, no_u_turns_within_tree = tf.while_loop(
          cond=lambda i, no_u_turn: (
              (i < read_indexes[1])
              & tf.reduce_any(  # pylint: disable=g-long-lambda
                  no_u_turn
              )
          ),
          body=_get_left_state_and_check_u_turn,
          loop_vars=(read_indexes[0], no_u_turns_within_tree),
      )
      return no_u_turns_within_tree
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 1 for calling a global function with example usage: has_not_u_turn_at_all_index(read_indexes, direction, momentum_state_memory, momentum_right, state_right, no_u_turns_within_tree, log_prob_rank, shard_axis_names) and returns: no_u_turns_within_tree, unknown,  unknown