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,
  ]
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 ) : 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.
Helper which processes input args to meet list - like assumptions .