code: def _prepare_args(
      target_log_prob_fn,
      volatility_fn,
      state,
      step_size,
      target_log_prob=None,
      grads_target_log_prob=None,
      volatility=None,
      grads_volatility_fn=None,
      diffusion_drift=None,
      parallel_iterations=10,
  ):
      """Helper which processes input args to meet list-like assumptions."""
      state_parts = (
          list(state) if mcmc_util.is_list_like(state) else [state]
      )
  
      [
          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,
      )
      [
          volatility_parts,
          grads_volatility,
      ] = _maybe_call_volatility_fn_and_grads(
          volatility_fn,
          state_parts,
          volatility,
          grads_volatility_fn,
          ps.shape(target_log_prob),
          parallel_iterations,
      )
  
      step_sizes = (
          list(step_size)
          if mcmc_util.is_list_like(step_size)
          else [step_size]
      )
      step_sizes = [
          tf.convert_to_tensor(
              value=s, name="step_size", dtype=target_log_prob.dtype
          )
          for s in step_sizes
      ]
      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`."
          )
  
      if diffusion_drift is None:
          diffusion_drift_parts = _get_drift(
              step_sizes,
              volatility_parts,
              grads_volatility,
              grads_target_log_prob,
          )
      else:
          diffusion_drift_parts = (
              list(diffusion_drift)
              if mcmc_util.is_list_like(diffusion_drift)
              else [diffusion_drift]
          )
          if len(state_parts) != len(diffusion_drift):
              raise ValueError(
                  "There should be exactly one `diffusion_drift` or it "
                  "should have same length as list-like `current_state`."
              )
  
      return [
          state_parts,
          step_sizes,
          target_log_prob,
          grads_target_log_prob,
          volatility_parts,
          grads_volatility,
          diffusion_drift_parts,
      ]
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 6 for calling a global function with example usage: _prepare_args(target_log_prob_fn, volatility_fn, state, step_size, target_log_prob, grads_target_log_prob, volatility, grads_volatility_fn, diffusion_drift, parallel_iterations) and returns: list