code: def _prepare_step_size(step_size, dtype, n_state_parts):
      step_sizes, _ = mcmc_util.prepare_state_parts(
          step_size, dtype=dtype, name="step_size"
      )
      if len(step_sizes) == 1:
          step_sizes *= n_state_parts
      if n_state_parts != len(step_sizes):
          raise ValueError(
              "There should be exactly one `step_size` or it should "
              "have same length as `current_state`."
          )
      return step_sizes
  
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 5 for calling a global function with example usage: _prepare_step_size(step_size, dtype, n_state_parts) and returns: step_sizes