tensorflow_probability/python/mcmc/replica_exchange_mc.py
def _get_exchanged_states(self, old_states, exchange_proposed,
                            exchange_proposed_n, sampled_replica_states,
                            sampled_replica_results):
    """Get list of TensorArrays holding exchanged states, and zeros."""
    with tf.compat.v1.name_scope('get_exchanged_states'):

      target_log_probs = []
      for replica in range(self.num_replica):
        replica_log_prob = _get_field(sampled_replica_results[replica],
                                      'target_log_prob')
        inverse_temp = self.inverse_temperatures[replica]
        target_log_probs.append(replica_log_prob / inverse_temp)
      target_log_probs = tf.stack(target_log_probs, axis=0)

      dtype = target_log_probs.dtype
      num_state_parts = len(sampled_replica_states[0])
      # exchanged_states[k][i] is Tensor of (new) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      # We will insert values into indices `i` for every replica with a proposed
      # exchange.
      exchanged_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              tensor_array_name='exchanged_states',
              # State part k has same shape, regardless of replica.  So use 0.
              element_shape=sampled_replica_states[0][k].shape)
          for k in range(num_state_parts)
      ]

      # Draw random variables here, to avoid sampling in the loop (and losing
      # reproducibility).  This may mean we sample too many, but we will always
      # have enough.
      sample_shape = tf.concat(
          ([self.num_replica // 2], tf.shape(input=target_log_probs)[1:]),
          axis=0)
      log_uniforms = tf.math.log(
          tf.random.uniform(
              shape=sample_shape, dtype=dtype, seed=self._seed_stream()))

      def _swap(is_exchange_accepted, x, y):
        """Swap batches of x, y where accepted."""
        with tf.compat.v1.name_scope('swap_where_exchange_accepted'):
          new_x = mcmc_util.choose(is_exchange_accepted, y, x)
          new_y = mcmc_util.choose(is_exchange_accepted, x, y)
        return new_x, new_y

      def cond(i, unused_exchanged_states):
        return i < exchange_proposed_n

      def body(i, exchanged_states):
        """Body of while loop for exchanging states."""
        # Propose exchange between replicas indexed by m and n.
        m, n = tf.unstack(exchange_proposed[i])

        # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff.
        # Note target_log_prob_diff = -EnergyDiff (common definition is in terms
        # of energy).
        temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]
        # Difference of target log probs may be +- Inf or NaN.  We want the
        # product of this with the temperature difference to have "alt value" of
        # -Inf.
        log_accept_ratio = mcmc_util.safe_sum(
            [-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])

        is_exchange_accepted = log_uniforms[i] < log_accept_ratio

        for k in range(num_state_parts):
          new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),
                               old_states[k].read(n))
          exchanged_states[k] = exchanged_states[k].write(m, new_m)
          exchanged_states[k] = exchanged_states[k].write(n, new_n)

        return i + 1, exchanged_states

      # At this point, exchanged_states[k] is a length num_replicas TensorArray.
      return tf.while_loop(
          cond=cond, body=body, loop_vars=[tf.constant(0),
                                           exchanged_states])[1]
def _get_exchanged_states(self, old_states, exchange_proposed,
                            exchange_proposed_n, sampled_replica_states,
                            sampled_replica_results):
    with tf.compat.v1.name_scope('get_exchanged_states'):

      target_log_probs = []
      for replica in range(self.num_replica):
        replica_log_prob = _get_field(sampled_replica_results[replica],
                                      'target_log_prob')
        inverse_temp = self.inverse_temperatures[replica]
        target_log_probs.append(replica_log_prob / inverse_temp)
      target_log_probs = tf.stack(target_log_probs, axis=0)

      dtype = target_log_probs.dtype
      num_state_parts = len(sampled_replica_states[0])
      # exchanged_states[k][i] is Tensor of (new) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      # We will insert values into indices `i` for every replica with a proposed
      # exchange.
      exchanged_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              tensor_array_name='exchanged_states',
              # State part k has same shape, regardless of replica.  So use 0.
              element_shape=sampled_replica_states[0][k].shape)
          for k in range(num_state_parts)
      ]

      # Draw random variables here, to avoid sampling in the loop (and losing
      # reproducibility).  This may mean we sample too many, but we will always
      # have enough.
      sample_shape = tf.concat(
          ([self.num_replica // 2], tf.shape(input=target_log_probs)[1:]),
          axis=0)
      log_uniforms = tf.math.log(
          tf.random.uniform(
              shape=sample_shape, dtype=dtype, seed=self._seed_stream()))

      def _swap(is_exchange_accepted, x, y):
        """Swap batches of x, y where accepted."""
        with tf.compat.v1.name_scope('swap_where_exchange_accepted'):
          new_x = mcmc_util.choose(is_exchange_accepted, y, x)
          new_y = mcmc_util.choose(is_exchange_accepted, x, y)
        return new_x, new_y

      def cond(i, unused_exchanged_states):
        return i < exchange_proposed_n

      def body(i, exchanged_states):
        """Body of while loop for exchanging states."""
        # Propose exchange between replicas indexed by m and n.
        m, n = tf.unstack(exchange_proposed[i])

        # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff.
        # Note target_log_prob_diff = -EnergyDiff (common definition is in terms
        # of energy).
        temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]
        # Difference of target log probs may be +- Inf or NaN.  We want the
        # product of this with the temperature difference to have "alt value" of
        # -Inf.
        log_accept_ratio = mcmc_util.safe_sum(
            [-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])

        is_exchange_accepted = log_uniforms[i] < log_accept_ratio

        for k in range(num_state_parts):
          new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),
                               old_states[k].read(n))
          exchanged_states[k] = exchanged_states[k].write(m, new_m)
          exchanged_states[k] = exchanged_states[k].write(n, new_n)

        return i + 1, exchanged_states

      # At this point, exchanged_states[k] is a length num_replicas TensorArray.
      return tf.while_loop(
          cond=cond, body=body, loop_vars=[tf.constant(0),
                                           exchanged_states])[1]
def _get_exchanged_states ( self , old_states , exchange_proposed , exchange_proposed_n , sampled_replica_states , sampled_replica_results ) : with tf . compat . v1 . name_scope ( 'get_exchanged_states' ) : target_log_probs = [ ] for replica in range ( self . num_replica ) : replica_log_prob = _get_field ( sampled_replica_results [ replica ] , 'target_log_prob' ) inverse_temp = self . inverse_temperatures [ replica ] target_log_probs . append ( replica_log_prob / inverse_temp ) target_log_probs = tf . stack ( target_log_probs , axis = 0 ) dtype = target_log_probs . dtype num_state_parts = len ( sampled_replica_states [ 0 ] ) # exchanged_states[k][i] is Tensor of (new) state part k, for replica i. # The `k` will be known statically, and `i` is a Tensor. # We will insert values into indices `i` for every replica with a proposed # exchange. exchanged_states = [ tf . TensorArray ( dtype , size = self . num_replica , dynamic_size = False , tensor_array_name = 'exchanged_states' , # State part k has same shape, regardless of replica.  So use 0. element_shape = sampled_replica_states [ 0 ] [ k ] . shape ) for k in range ( num_state_parts ) ] # Draw random variables here, to avoid sampling in the loop (and losing # reproducibility).  This may mean we sample too many, but we will always # have enough. sample_shape = tf . concat ( ( [ self . num_replica // 2 ] , tf . shape ( input = target_log_probs ) [ 1 : ] ) , axis = 0 ) log_uniforms = tf . math . log ( tf . random . uniform ( shape = sample_shape , dtype = dtype , seed = self . _seed_stream ( ) ) ) def _swap ( is_exchange_accepted , x , y ) : """Swap batches of x, y where accepted.""" with tf . compat . v1 . name_scope ( 'swap_where_exchange_accepted' ) : new_x = mcmc_util . choose ( is_exchange_accepted , y , x ) new_y = mcmc_util . choose ( is_exchange_accepted , x , y ) return new_x , new_y def cond ( i , unused_exchanged_states ) : return i < exchange_proposed_n def body ( i , exchanged_states ) : """Body of while loop for exchanging states.""" # Propose exchange between replicas indexed by m and n. m , n = tf . unstack ( exchange_proposed [ i ] ) # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff. # Note target_log_prob_diff = -EnergyDiff (common definition is in terms # of energy). temp_diff = self . inverse_temperatures [ m ] - self . inverse_temperatures [ n ] # Difference of target log probs may be +- Inf or NaN.  We want the # product of this with the temperature difference to have "alt value" of # -Inf. log_accept_ratio = mcmc_util . safe_sum ( [ - temp_diff * target_log_probs [ m ] , temp_diff * target_log_probs [ n ] ] ) is_exchange_accepted = log_uniforms [ i ] < log_accept_ratio for k in range ( num_state_parts ) : new_m , new_n = _swap ( is_exchange_accepted , old_states [ k ] . read ( m ) , old_states [ k ] . read ( n ) ) exchanged_states [ k ] = exchanged_states [ k ] . write ( m , new_m ) exchanged_states [ k ] = exchanged_states [ k ] . write ( n , new_n ) return i + 1 , exchanged_states # At this point, exchanged_states[k] is a length num_replicas TensorArray. return tf . while_loop ( cond = cond , body = body , loop_vars = [ tf . constant ( 0 ) , exchanged_states ] ) [ 1 ]
Get list of TensorArrays holding exchanged states, and zeros.
Get list of TensorArrays holding exchanged states and zeros .