tensorflow_probability/python/mcmc/langevin.py
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,
      distribution_util.prefer_static_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,
  ]
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,
      distribution_util.prefer_static_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,
  ]
Helper which processes input args to meet list-like assumptions.