class_docstr: `Reducer` that accumulates trace results at each sample.
  
  Trace results are defined by an appropriate `trace_fn`, which accepts the
  current chain state and kernel results, and returns the desired result.
  At each sample, the traced values are added to a `TensorArray` which
  accumulates all results. By default, all kernel results are traced but in the
  future the default will be changed to no results being traced, so plan
  accordingly.
  
  If wrapped in a `tfp.experimental.mcmc.WithReductions` Transition Kernel,
  TracingReducer will not accumulate the kernel results of `WithReductions`.
  Rather, the top level kernel results will be that of `WithReductions`' inner
  kernel.
  
  As with all reducers, `TracingReducer` does not hold state information;
  rather, it stores supplied metadata. Intermediate calculations are held in
  a `TracingState` named tuple, which is returned via `initialize` and
  `one_step` method calls.
class_docstr_tok: ["`","Reducer","`","that","accumulates","trace","results","at","each","sample",".","Trace","results","are","defined","by","an","appropriate","`","trace_fn","`",",","which","accepts","the","current","chain","state","and","kernel","results",",","and","returns","the","desired","result",".","At","each","sample",",","the","traced","values","are","added","to","a","`","TensorArray","`","which","accumulates","all","results",".","By","default",",","all","kernel","results","are","traced","but","in","the","future","the","default","will","be","changed","to","no","results","being","traced",",","so","plan","accordingly",".","If","wrapped","in","a","`","tfp.experimental.mcmc.WithReductions","`","Transition","Kernel",",","TracingReducer","will","not","accumulate","the","kernel","results","of","`","WithReductions","`",".","Rather",",","the","top","level","kernel","results","will","be","that","of","`","WithReductions","`","'","inner","kernel",".","As","with","all","reducers",",","`","TracingReducer","`","does","not","hold","state","information",";","rather",",","it","stores","supplied","metadata",".","Intermediate","calculations","are","held","in","a","`","TracingState","`","named","tuple",",","which","is","returned","via","`","initialize","`","and","`","one_step","`","method","calls","."]
docstr: Update the `current_reducer_state` with a new trace result.
  
  The trace result will be computed by evaluating the `trace_fn` provided
  at instantiation with the `new_chain_state` and `previous_kernel_results`.
  
  Args:
    new_chain_state: A (possibly nested) structure of incoming chain state(s)
      with shape and dtype compatible with those used to initialize the
      `TracingState`.
    current_reducer_state: `TracingState`s representing all previously traced
      results.
    previous_kernel_results: A (possibly nested) structure of `Tensor`s
      representing internal calculations made in a related
      `TransitionKernel`.
  
  Returns:
    new_reducer_state: `TracingState` with updated trace. Its `trace_state`
      field holds a `TensorArray` that includes the newly computed trace
      result.
docstr_tok: ["Update","the","`","current_reducer_state","`","with","a","new","trace","result",".","The","trace","result","will","be","computed","by","evaluating","the","`","trace_fn","`","provided","at","instantiation","with","the","`","new_chain_state","`","and","`","previous_kernel_results","`",".","Args",":","new_chain_state",":","A","(","possibly","nested",")","structure","of","incoming","chain","state","(","s",")","with","shape","and","dtype","compatible","with","those","used","to","initialize","the","`","TracingState","`",".","current_reducer_state",":","`","TracingState","`","s","representing","all","previously","traced","results",".","previous_kernel_results",":","A","(","possibly","nested",")","structure","of","`","Tensor","`","s","representing","internal","calculations","made","in","a","related","`","TransitionKernel","`",".","Returns",":","new_reducer_state",":","`","TracingState","`","with","updated","trace",".","Its","`","trace_state","`","field","holds","a","`","TensorArray","`","that","includes","the","newly","computed","trace","result","."]
code: def one_step(
      self,
      new_chain_state,
      current_reducer_state,
      previous_kernel_results,
  ):
      """Update the `current_reducer_state` with a new trace result.
  
      The trace result will be computed by evaluating the `trace_fn` provided
      at instantiation with the `new_chain_state` and `previous_kernel_results`.
  
      Args:
        new_chain_state: A (possibly nested) structure of incoming chain state(s)
          with shape and dtype compatible with those used to initialize the
          `TracingState`.
        current_reducer_state: `TracingState`s representing all previously traced
          results.
        previous_kernel_results: A (possibly nested) structure of `Tensor`s
          representing internal calculations made in a related
          `TransitionKernel`.
  
      Returns:
        new_reducer_state: `TracingState` with updated trace. Its `trace_state`
          field holds a `TensorArray` that includes the newly computed trace
          result.
      """
      with tf.name_scope(
          mcmc_util.make_name(self.name, "tracing_reducer", "one_step")
      ):
          num_samples, trace_state = current_reducer_state
          new_trace_results = self.trace_fn(
              new_chain_state, previous_kernel_results
          )
          new_trace_state = tf.nest.map_structure(
              lambda state, new_trace: state.write(
                  num_samples, new_trace
              ),
              trace_state,
              new_trace_results,
          )
          return TracingState(num_samples + 1, new_trace_state)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/tracing_reducer.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import reducer
  from tensorflow_probability.python.internal import loop_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 3 for calling the TracingReducer obj's underlying member method code with example usage: obj.one_step(new_chain_state, current_reducer_state, previous_kernel_results) and returns: TracingState