tensorflow_probability/python/mcmc/replica_exchange_mc.py
def one_step(self, current_state, previous_kernel_results):
    """Takes one step of the TransitionKernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s).
      previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
        `list` of `Tensor`s representing internal calculations made within the
        previous call to this function (or as returned by `bootstrap_results`).

    Returns:
      next_state: `Tensor` or Python `list` of `Tensor`s representing the
        next state(s) of the Markov chain(s).
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within this function.
        This inculdes replica states.
    """
    # Key difficulty:  The type of exchanges differs from one call to the
    # next...even the number of exchanges can differ.
    # As a result, exchanges must happen dynamically, in while loops.
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'remc', 'one_step'),
        values=[current_state, previous_kernel_results]):

      # Each replica does `one_step` to get pre-exchange states/KernelResults.
      sampled_replica_states, sampled_replica_results = zip(*[
          rk.one_step(previous_kernel_results.replica_states[i],
                      previous_kernel_results.replica_results[i])
          for i, rk in enumerate(self.replica_kernels)
      ])
      sampled_replica_states = list(sampled_replica_states)
      sampled_replica_results = list(sampled_replica_results)

      states_are_lists = mcmc_util.is_list_like(sampled_replica_states[0])

      if not states_are_lists:
        sampled_replica_states = [[s] for s in sampled_replica_states]
      num_state_parts = len(sampled_replica_states[0])

      dtype = sampled_replica_states[0][0].dtype

      # Must put states into TensorArrays.  Why?  We will read/write states
      # dynamically with Tensor index `i`, and you cannot do this with lists.
      # old_states[k][i] is Tensor of (old) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      old_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              clear_after_read=False,
              tensor_array_name='old_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)
      ]
      for k in range(num_state_parts):
        for i in range(self.num_replica):
          old_states[k] = old_states[k].write(i, sampled_replica_states[i][k])

      exchange_proposed = self.exchange_proposed_fn(
          self.num_replica, seed=self._seed_stream())
      exchange_proposed_n = tf.shape(input=exchange_proposed)[0]

      exchanged_states = self._get_exchanged_states(
          old_states, exchange_proposed, exchange_proposed_n,
          sampled_replica_states, sampled_replica_results)

      no_exchange_proposed, _ = tf.compat.v1.setdiff1d(
          tf.range(self.num_replica), tf.reshape(exchange_proposed, [-1]))

      exchanged_states = self._insert_old_states_where_no_exchange_was_proposed(
          no_exchange_proposed, old_states, exchanged_states)

      next_replica_states = []
      for i in range(self.num_replica):
        next_replica_states_i = []
        for k in range(num_state_parts):
          next_replica_states_i.append(exchanged_states[k].read(i))
        next_replica_states.append(next_replica_states_i)

      if not states_are_lists:
        next_replica_states = [s[0] for s in next_replica_states]
        sampled_replica_states = [s[0] for s in sampled_replica_states]

      # Now that states are/aren't exchanged, bootstrap next kernel_results.
      # The viewpoint is that after each exchange, we are starting anew.
      next_replica_results = [
          rk.bootstrap_results(state)
          for rk, state in zip(self.replica_kernels, next_replica_states)
      ]

      next_state = next_replica_states[0]  # Replica 0 is the returned state(s).

      kernel_results = ReplicaExchangeMCKernelResults(
          replica_states=next_replica_states,
          replica_results=next_replica_results,
          sampled_replica_states=sampled_replica_states,
          sampled_replica_results=sampled_replica_results,
      )

      return next_state, kernel_results
def one_step(self, current_state, previous_kernel_results):
    """Takes one step of the TransitionKernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s).
      previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
        `list` of `Tensor`s representing internal calculations made within the
        previous call to this function (or as returned by `bootstrap_results`).

    Returns:
      next_state: `Tensor` or Python `list` of `Tensor`s representing the
        next state(s) of the Markov chain(s).
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within this function.
        This inculdes replica states.
    """
    # Key difficulty:  The type of exchanges differs from one call to the
    # next...even the number of exchanges can differ.
    # As a result, exchanges must happen dynamically, in while loops.
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'remc', 'one_step'),
        values=[current_state, previous_kernel_results]):

      # Each replica does `one_step` to get pre-exchange states/KernelResults.
      sampled_replica_states, sampled_replica_results = zip(*[
          rk.one_step(previous_kernel_results.replica_states[i],
                      previous_kernel_results.replica_results[i])
          for i, rk in enumerate(self.replica_kernels)
      ])
      sampled_replica_states = list(sampled_replica_states)
      sampled_replica_results = list(sampled_replica_results)

      states_are_lists = mcmc_util.is_list_like(sampled_replica_states[0])

      if not states_are_lists:
        sampled_replica_states = [[s] for s in sampled_replica_states]
      num_state_parts = len(sampled_replica_states[0])

      dtype = sampled_replica_states[0][0].dtype

      # Must put states into TensorArrays.  Why?  We will read/write states
      # dynamically with Tensor index `i`, and you cannot do this with lists.
      # old_states[k][i] is Tensor of (old) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      old_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              clear_after_read=False,
              tensor_array_name='old_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)
      ]
      for k in range(num_state_parts):
        for i in range(self.num_replica):
          old_states[k] = old_states[k].write(i, sampled_replica_states[i][k])

      exchange_proposed = self.exchange_proposed_fn(
          self.num_replica, seed=self._seed_stream())
      exchange_proposed_n = tf.shape(input=exchange_proposed)[0]

      exchanged_states = self._get_exchanged_states(
          old_states, exchange_proposed, exchange_proposed_n,
          sampled_replica_states, sampled_replica_results)

      no_exchange_proposed, _ = tf.compat.v1.setdiff1d(
          tf.range(self.num_replica), tf.reshape(exchange_proposed, [-1]))

      exchanged_states = self._insert_old_states_where_no_exchange_was_proposed(
          no_exchange_proposed, old_states, exchanged_states)

      next_replica_states = []
      for i in range(self.num_replica):
        next_replica_states_i = []
        for k in range(num_state_parts):
          next_replica_states_i.append(exchanged_states[k].read(i))
        next_replica_states.append(next_replica_states_i)

      if not states_are_lists:
        next_replica_states = [s[0] for s in next_replica_states]
        sampled_replica_states = [s[0] for s in sampled_replica_states]

      # Now that states are/aren't exchanged, bootstrap next kernel_results.
      # The viewpoint is that after each exchange, we are starting anew.
      next_replica_results = [
          rk.bootstrap_results(state)
          for rk, state in zip(self.replica_kernels, next_replica_states)
      ]

      next_state = next_replica_states[0]  # Replica 0 is the returned state(s).

      kernel_results = ReplicaExchangeMCKernelResults(
          replica_states=next_replica_states,
          replica_results=next_replica_results,
          sampled_replica_states=sampled_replica_states,
          sampled_replica_results=sampled_replica_results,
      )

      return next_state, kernel_results
def one_step ( self , current_state , previous_kernel_results ) : # Key difficulty:  The type of exchanges differs from one call to the # next...even the number of exchanges can differ. # As a result, exchanges must happen dynamically, in while loops. with tf . compat . v1 . name_scope ( name = mcmc_util . make_name ( self . name , 'remc' , 'one_step' ) , values = [ current_state , previous_kernel_results ] ) : # Each replica does `one_step` to get pre-exchange states/KernelResults. sampled_replica_states , sampled_replica_results = zip ( * [ rk . one_step ( previous_kernel_results . replica_states [ i ] , previous_kernel_results . replica_results [ i ] ) for i , rk in enumerate ( self . replica_kernels ) ] ) sampled_replica_states = list ( sampled_replica_states ) sampled_replica_results = list ( sampled_replica_results ) states_are_lists = mcmc_util . is_list_like ( sampled_replica_states [ 0 ] ) if not states_are_lists : sampled_replica_states = [ [ s ] for s in sampled_replica_states ] num_state_parts = len ( sampled_replica_states [ 0 ] ) dtype = sampled_replica_states [ 0 ] [ 0 ] . dtype # Must put states into TensorArrays.  Why?  We will read/write states # dynamically with Tensor index `i`, and you cannot do this with lists. # old_states[k][i] is Tensor of (old) state part k, for replica i. # The `k` will be known statically, and `i` is a Tensor. old_states = [ tf . TensorArray ( dtype , size = self . num_replica , dynamic_size = False , clear_after_read = False , tensor_array_name = 'old_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 ) ] for k in range ( num_state_parts ) : for i in range ( self . num_replica ) : old_states [ k ] = old_states [ k ] . write ( i , sampled_replica_states [ i ] [ k ] ) exchange_proposed = self . exchange_proposed_fn ( self . num_replica , seed = self . _seed_stream ( ) ) exchange_proposed_n = tf . shape ( input = exchange_proposed ) [ 0 ] exchanged_states = self . _get_exchanged_states ( old_states , exchange_proposed , exchange_proposed_n , sampled_replica_states , sampled_replica_results ) no_exchange_proposed , _ = tf . compat . v1 . setdiff1d ( tf . range ( self . num_replica ) , tf . reshape ( exchange_proposed , [ - 1 ] ) ) exchanged_states = self . _insert_old_states_where_no_exchange_was_proposed ( no_exchange_proposed , old_states , exchanged_states ) next_replica_states = [ ] for i in range ( self . num_replica ) : next_replica_states_i = [ ] for k in range ( num_state_parts ) : next_replica_states_i . append ( exchanged_states [ k ] . read ( i ) ) next_replica_states . append ( next_replica_states_i ) if not states_are_lists : next_replica_states = [ s [ 0 ] for s in next_replica_states ] sampled_replica_states = [ s [ 0 ] for s in sampled_replica_states ] # Now that states are/aren't exchanged, bootstrap next kernel_results. # The viewpoint is that after each exchange, we are starting anew. next_replica_results = [ rk . bootstrap_results ( state ) for rk , state in zip ( self . replica_kernels , next_replica_states ) ] next_state = next_replica_states [ 0 ] # Replica 0 is the returned state(s). kernel_results = ReplicaExchangeMCKernelResults ( replica_states = next_replica_states , replica_results = next_replica_results , sampled_replica_states = sampled_replica_states , sampled_replica_results = sampled_replica_results , ) return next_state , kernel_results
Takes one step of the TransitionKernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s).
      previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
        `list` of `Tensor`s representing internal calculations made within the
        previous call to this function (or as returned by `bootstrap_results`).

    Returns:
      next_state: `Tensor` or Python `list` of `Tensor`s representing the
        next state(s) of the Markov chain(s).
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within this function.
        This inculdes replica states.
Takes one step of the TransitionKernel .