experimental/neutra/neutra_kernel.py
def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
      return kernel_results._replace(
          accepted_results=kernel_results.accepted_results._replace(
              num_leapfrog_steps=num_leapfrog_steps))

    step_size = previous_kernel_results.new_step_size
    previous_kernel_results = set_num_leapfrog_steps(
        previous_kernel_results, self._num_leapfrog_steps(step_size))

    new_state, kernel_results = self._kernel.one_step(
        self._flatten_state(current_state), previous_kernel_results)
    return self._unflatten_state(new_state), kernel_results
def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
      return kernel_results._replace(
          accepted_results=kernel_results.accepted_results._replace(
              num_leapfrog_steps=num_leapfrog_steps))

    step_size = previous_kernel_results.new_step_size
    previous_kernel_results = set_num_leapfrog_steps(
        previous_kernel_results, self._num_leapfrog_steps(step_size))

    new_state, kernel_results = self._kernel.one_step(
        self._flatten_state(current_state), previous_kernel_results)
    return self._unflatten_state(new_state), kernel_results
Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.