experimental/fun_mcmc/fun_mcmc_lib.py
def metropolis_hastings_step(current_state: State,
                             proposed_state: State,
                             energy_change: FloatTensor,
                             seed=None) -> Tuple[State, tf.Tensor, tf.Tensor]:
  """Metropolis-Hastings step.

  This probabilistically chooses between `current_state` and `proposed_state`
  based on the `energy_change` so as to preserve detailed balance.

  Energy change is the negative of `log_accept_ratio`.

  Args:
    current_state: Current state.
    proposed_state: Proposed state.
    energy_change: E(proposed_state) - E(previous_state).
    seed: For reproducibility.

  Returns:
    new_state: The chosen state.
    is_accepted: Whether the proposed state was accepted.
    log_uniform: The random number that was used to select between the two
      states.
  """
  flat_current = tf.nest.flatten(current_state)
  flat_proposed = nest.flatten_up_to(current_state, proposed_state)
  # Impute the None's in the current state.
  flat_current = [
      p if c is None else c for p, c in zip(flat_proposed, flat_current)
  ]
  current_state = tf.nest.pack_sequence_as(current_state, flat_current)

  current_state = tf.nest.map_structure(tf.convert_to_tensor, current_state)
  proposed_state = tf.nest.map_structure(tf.convert_to_tensor, proposed_state)
  energy_change = tf.convert_to_tensor(value=energy_change)

  log_accept_ratio = -energy_change

  log_uniform = tf.math.log(
      tf.random.uniform(
          shape=tf.shape(input=log_accept_ratio),
          dtype=log_accept_ratio.dtype.base_dtype,
          seed=seed))
  is_accepted = log_uniform < log_accept_ratio

  next_state = mcmc_util.choose(
      is_accepted, proposed_state, current_state, name='choose_next_state')
  return next_state, is_accepted, log_uniform
def metropolis_hastings_step(current_state: State,
                             proposed_state: State,
                             energy_change: FloatTensor,
                             seed=None) -> Tuple[State, tf.Tensor, tf.Tensor]:
  flat_current = tf.nest.flatten(current_state)
  flat_proposed = nest.flatten_up_to(current_state, proposed_state)
  # Impute the None's in the current state.
  flat_current = [
      p if c is None else c for p, c in zip(flat_proposed, flat_current)
  ]
  current_state = tf.nest.pack_sequence_as(current_state, flat_current)

  current_state = tf.nest.map_structure(tf.convert_to_tensor, current_state)
  proposed_state = tf.nest.map_structure(tf.convert_to_tensor, proposed_state)
  energy_change = tf.convert_to_tensor(value=energy_change)

  log_accept_ratio = -energy_change

  log_uniform = tf.math.log(
      tf.random.uniform(
          shape=tf.shape(input=log_accept_ratio),
          dtype=log_accept_ratio.dtype.base_dtype,
          seed=seed))
  is_accepted = log_uniform < log_accept_ratio

  next_state = mcmc_util.choose(
      is_accepted, proposed_state, current_state, name='choose_next_state')
  return next_state, is_accepted, log_uniform
def metropolis_hastings_step ( current_state : State , proposed_state : State , energy_change : FloatTensor , seed = None ) -> Tuple [ State , tf . Tensor , tf . Tensor ] : flat_current = tf . nest . flatten ( current_state ) flat_proposed = nest . flatten_up_to ( current_state , proposed_state ) # Impute the None's in the current state. flat_current = [ p if c is None else c for p , c in zip ( flat_proposed , flat_current ) ] current_state = tf . nest . pack_sequence_as ( current_state , flat_current ) current_state = tf . nest . map_structure ( tf . convert_to_tensor , current_state ) proposed_state = tf . nest . map_structure ( tf . convert_to_tensor , proposed_state ) energy_change = tf . convert_to_tensor ( value = energy_change ) log_accept_ratio = - energy_change log_uniform = tf . math . log ( tf . random . uniform ( shape = tf . shape ( input = log_accept_ratio ) , dtype = log_accept_ratio . dtype . base_dtype , seed = seed ) ) is_accepted = log_uniform < log_accept_ratio next_state = mcmc_util . choose ( is_accepted , proposed_state , current_state , name = 'choose_next_state' ) return next_state , is_accepted , log_uniform
Metropolis-Hastings step.

  This probabilistically chooses between `current_state` and `proposed_state`
  based on the `energy_change` so as to preserve detailed balance.

  Energy change is the negative of `log_accept_ratio`.

  Args:
    current_state: Current state.
    proposed_state: Proposed state.
    energy_change: E(proposed_state) - E(previous_state).
    seed: For reproducibility.

  Returns:
    new_state: The chosen state.
    is_accepted: Whether the proposed state was accepted.
    log_uniform: The random number that was used to select between the two
      states.
Metropolis - Hastings step .