code: def _ensure_step_size_is_scalar(step_size, validate_args):
      """Ensures that the step size is a scalar `Tensor`."""
      if tf.nest.is_nested(step_size):
          raise ValueError(
              "Step size must be a scalar. Got: {}".format(step_size)
          )
      rank = ps.rank(step_size)
      rank_ = tf.get_static_value(rank)
      if rank_ is not None:
          if rank_ != 0:
              raise ValueError(
                  "Step size must be a scalar. Got: {}".format(
                      step_size
                  )
              )
      elif validate_args:
          with tf.control_dependencies(
              [
                  assert_util.assert_rank(
                      step_size, 0, "Step size must be a scalar."
                  )
              ]
          ):
              return tf.identity(step_size)
      return step_size
  
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 13 for calling a global function with example usage: _ensure_step_size_is_scalar(step_size, validate_args) and returns: step_size, tf