tensorflow_probability/python/mcmc/hmc.py
def _prepare_args(target_log_prob_fn,
                  state,
                  step_size,
                  target_log_prob=None,
                  grads_target_log_prob=None,
                  maybe_expand=False,
                  state_gradients_are_stopped=False):
  """Helper which processes input args to meet list-like assumptions."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  state_parts = [
      tf.convert_to_tensor(value=s, name='current_state') for s in state_parts
  ]
  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 = (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`.')
  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),
      target_log_prob,
      grads_target_log_prob,
  ]
def _prepare_args(target_log_prob_fn,
                  state,
                  step_size,
                  target_log_prob=None,
                  grads_target_log_prob=None,
                  maybe_expand=False,
                  state_gradients_are_stopped=False):
  """Helper which processes input args to meet list-like assumptions."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  state_parts = [
      tf.convert_to_tensor(value=s, name='current_state') for s in state_parts
  ]
  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 = (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`.')
  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),
      target_log_prob,
      grads_target_log_prob,
  ]
def _prepare_args ( target_log_prob_fn , state , step_size , target_log_prob = None , grads_target_log_prob = None , maybe_expand = False , state_gradients_are_stopped = False ) : state_parts = list ( state ) if mcmc_util . is_list_like ( state ) else [ state ] state_parts = [ tf . convert_to_tensor ( value = s , name = 'current_state' ) for s in state_parts ] 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 = ( 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`.' ) 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 ) , target_log_prob , grads_target_log_prob , ]
Helper which processes input args to meet list-like assumptions.
Helper which processes input args to meet list - like assumptions .