class_docstr: Applies `Reducer`s to stream over MCMC samples.
  
  `WithReductions` augments an inner MCMC kernel with
  side-computations that can read the stream of samples as they
  are generated.  This is relevant for streaming uses of MCMC,
  where materializing the entire Markov chain history is undesirable,
  e.g. due to memory limits.
  
  One `WithReductions` instance can attach an arbitrary collection
  of side-computations, each of which must be packaged as a
  `Reducer`.  `WithReductions` operates by generating a
  sample with its `inner_kernel`'s `one_step`, then invoking each
  `Reducer`'s `one_step` method on that sample. The updated reducer
  states are stored in the `reduction_results` field of
  `WithReductions`' kernel results.
class_docstr_tok: ["Applies","`","Reducer","`","s","to","stream","over","MCMC","samples",".","`","WithReductions","`","augments","an","inner","MCMC","kernel","with","side-computations","that","can","read","the","stream","of","samples","as","they","are","generated",".","This","is","relevant","for","streaming","uses","of","MCMC",",","where","materializing","the","entire","Markov","chain","history","is","undesirable",",","e.g",".","due","to","memory","limits",".","One","`","WithReductions","`","instance","can","attach","an","arbitrary","collection","of","side-computations",",","each","of","which","must","be","packaged","as","a","`","Reducer","`",".","`","WithReductions","`","operates","by","generating","a","sample","with","its","`","inner_kernel","`","'s","`","one_step","`",",","then","invoking","each","`","Reducer","`","'s","`","one_step","`","method","on","that","sample",".","The","updated","reducer","states","are","stored","in","the","`","reduction_results","`","field","of","`","WithReductions","`","'","kernel","results","."]
docstr: Updates all `Reducer`s with a new sample from the `inner_kernel`.
  
  Args:
    current_state: `Tensor` or Python `list` of `Tensor`s
      representing the current state(s) of the Markov chain(s),
    previous_kernel_results: `WithReductionsKernelResults` named tuple.
      `WithReductionsKernelResults` contain the state of
      `reduction_results` and a reference to kernel results of
      nested `TransitionKernel`s.
    seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
  Returns:
    new_state: Newest MCMC state drawn from the `inner_kernel`.
    kernel_results: `WithReductionsKernelResults` representing updated
      kernel results. Reducer states are stored in the
      `reduction_results` field. The state structure is identical
      to `self.reducer`.
docstr_tok: ["Updates","all","`","Reducer","`","s","with","a","new","sample","from","the","`","inner_kernel","`",".","Args",":","current_state",":","`","Tensor","`","or","Python","`","list","`","of","`","Tensor","`","s","representing","the","current","state","(","s",")","of","the","Markov","chain","(","s",")",",","previous_kernel_results",":","`","WithReductionsKernelResults","`","named","tuple",".","`","WithReductionsKernelResults","`","contain","the","state","of","`","reduction_results","`","and","a","reference","to","kernel","results","of","nested","`","TransitionKernel","`","s",".","seed",":","PRNG","seed",";","see","`","tfp.random.sanitize_seed","`","for","details",".","Returns",":","new_state",":","Newest","MCMC","state","drawn","from","the","`","inner_kernel","`",".","kernel_results",":","`","WithReductionsKernelResults","`","representing","updated","kernel","results",".","Reducer","states","are","stored","in","the","`","reduction_results","`","field",".","The","state","structure","is","identical","to","`","self.reducer","`","."]
code: def one_step(self, current_state, previous_kernel_results, seed=None):
      """Updates all `Reducer`s with a new sample from the `inner_kernel`.
  
      Args:
        current_state: `Tensor` or Python `list` of `Tensor`s
          representing the current state(s) of the Markov chain(s),
        previous_kernel_results: `WithReductionsKernelResults` named tuple.
          `WithReductionsKernelResults` contain the state of
          `reduction_results` and a reference to kernel results of
          nested `TransitionKernel`s.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        new_state: Newest MCMC state drawn from the `inner_kernel`.
        kernel_results: `WithReductionsKernelResults` representing updated
          kernel results. Reducer states are stored in the
          `reduction_results` field. The state structure is identical
          to `self.reducer`.
      """
      with tf.name_scope(
          mcmc_util.make_name(self.name, "with_reductions", "one_step")
      ):
          new_state, inner_kernel_results = self.inner_kernel.one_step(
              current_state,
              previous_kernel_results.inner_results,
              seed=seed,
          )
          inner_kernel_results_adj = self.adjust_kr_fn(
              inner_kernel_results
          )
  
          def step_reducer(r, state):
              return r.one_step(
                  new_state,
                  state,
                  previous_kernel_results=inner_kernel_results_adj,
              )
  
          new_reducer_state = nest.map_structure_up_to(
              self.reducer,
              step_reducer,
              self.reducer,
              previous_kernel_results.reduction_results,
              check_types=False,
          )
          kernel_results = WithReductionsKernelResults(
              new_reducer_state, inner_kernel_results
          )
          return new_state, kernel_results
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/with_reductions.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling the WithReductions obj's underlying member method code with example usage: obj.one_step(current_state, previous_kernel_results, seed) and returns: new_state,  kernel_results, r