code: def _prepare_args(
      target_log_prob_fn,
      state,
      step_size,
      momentum_distribution,
      target_log_prob=None,
      grads_target_log_prob=None,
      maybe_expand=False,
      state_gradients_are_stopped=False,
      experimental_shard_axis_names=None,
  ):
      """Helper which processes input args to meet list-like assumptions."""
      state_parts, _ = mcmc_util.prepare_state_parts(
          state, name="current_state"
      )
      if state_gradients_are_stopped:
          state_parts = [tf.stop_gradient(x) for x in state_parts]
      (
          target_log_prob,
          grads_target_log_prob,
      ) = mcmc_util.maybe_call_fn_and_grads(
          target_log_prob_fn,
          state_parts,
          target_log_prob,
          grads_target_log_prob,
      )
      step_sizes, _ = mcmc_util.prepare_state_parts(
          step_size, dtype=target_log_prob.dtype, name="step_size"
      )
  
      # Default momentum distribution is None
      if momentum_distribution is None:
          momentum_distribution = pu.make_momentum_distribution(
              state_parts,
              ps.shape(target_log_prob),
              shard_axis_names=experimental_shard_axis_names,
          )
      momentum_distribution = pu.maybe_make_list_and_batch_broadcast(
          momentum_distribution, ps.shape(target_log_prob)
      )
  
      if len(step_sizes) == 1:
          step_sizes *= len(state_parts)
      if len(state_parts) != len(step_sizes):
          raise ValueError(
              "There should be exactly one `step_size` or it should "
              "have same length as `current_state`."
          )
  
      def maybe_flatten(x):
          return (
              x
              if maybe_expand or mcmc_util.is_list_like(state)
              else x[0]
          )
  
      return [
          maybe_flatten(state_parts),
          maybe_flatten(step_sizes),
          momentum_distribution,
          target_log_prob,
          grads_target_log_prob,
      ]
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import metropolis_hastings
  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: _prepare_args(target_log_prob_fn, state, step_size, momentum_distribution, target_log_prob, grads_target_log_prob, maybe_expand, state_gradients_are_stopped, experimental_shard_axis_names) and returns: list, unknown