code: def _check_at_least_two_chains(
      accept_prob, reduce_chain_axis_names, validate_args, message
  ):
      """Checks that the number of chains is at least 2."""
      # Number of total chains is local batch size * distributed axis size
      local_axis_size = ps.size(accept_prob)
      distributed_axis_size = int(
          ps.reduce_prod(
              [
                  distribute_lib.get_axis_size(a)
                  for a in reduce_chain_axis_names
              ]
          )
      )
      num_chains = local_axis_size * distributed_axis_size
      num_chains_ = tf.get_static_value(num_chains)
      if num_chains_ is not None:
          if num_chains_ < 2:
              raise ValueError(
                  "{} Got: {}".format(message, num_chains_)
              )
      elif validate_args:
          with tf.control_dependencies(
              [assert_util.assert_greater_equal(num_chains, 2, message)]
          ):
              accept_prob = tf.identity(accept_prob)
      return accept_prob
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  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 unnest
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 17 for calling a global function with example usage: _check_at_least_two_chains(accept_prob, reduce_chain_axis_names, validate_args, message) and returns: accept_prob