code: def has_not_u_turn(
      state_diff,
      momentum_left,
      momentum_right,
      log_prob_rank,
      shard_axis_names=None,
  ):
      """If the trajectory does not exhibit a U-turn pattern."""
      shard_axis_names = shard_axis_names or ([None] * len(state_diff))
  
      def reduce_sum(x, m, shard_axes):
          out = tf.reduce_sum(
              x, axis=ps.range(log_prob_rank, ps.rank(m))
          )
          if shard_axes is not None:
              out = distribute_lib.psum(out, shard_axes)
          return out
  
      with tf.name_scope("has_not_u_turn"):
          batch_dot_product_left = sum(
              reduce_sum(s_diff * m, m, axes)
              for s_diff, m, axes in zip(
                  state_diff, momentum_left, shard_axis_names
              )
          )
          batch_dot_product_right = sum(
              reduce_sum(s_diff * m, m, axes)
              for s_diff, m, axes in zip(
                  state_diff, momentum_right, shard_axis_names
              )
          )
          return (batch_dot_product_left >= 0) & (
              batch_dot_product_right >= 0
          )
  
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 2 for calling a global function with example usage: has_not_u_turn(state_diff, momentum_left, momentum_right, log_prob_rank, shard_axis_names) and returns: out, unknown