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: Initializes a `TracingState` using previously defined metadata.
  
  Both the `initial_chain_state` and `initial_kernel_results` do not count
  as a sample, and hence, will not be traced. This is a deliberate decision
  that ensures consistency across sampling procedures.
  
  Args:
    initial_chain_state: A (possibly nested) structure of `Tensor`s or Python
      `list`s of `Tensor`s representing the current state(s) of the Markov
      chain(s). It is used to infer the structure of future trace results.
    initial_kernel_results: A (possibly nested) structure of `Tensor`s
      representing internal calculations made in a related `TransitionKernel`.
      It is used to infer the structure of future trace results.
  
  Returns:
    state: `TracingState` with an empty `TensorArray` in its `trace_state`
      field.
docstr_tok: ["Initializes","a","`","TracingState","`","using","previously","defined","metadata",".","Both","the","`","initial_chain_state","`","and","`","initial_kernel_results","`","do","not","count","as","a","sample",",","and","hence",",","will","not","be","traced",".","This","is","a","deliberate","decision","that","ensures","consistency","across","sampling","procedures",".","Args",":","initial_chain_state",":","A","(","possibly","nested",")","structure","of","`","Tensor","`","s","or","Python","`","list","`","s","of","`","Tensor","`","s","representing","the","current","state","(","s",")","of","the","Markov","chain","(","s",")",".","It","is","used","to","infer","the","structure","of","future","trace","results",".","initial_kernel_results",":","A","(","possibly","nested",")","structure","of","`","Tensor","`","s","representing","internal","calculations","made","in","a","related","`","TransitionKernel","`",".","It","is","used","to","infer","the","structure","of","future","trace","results",".","Returns",":","state",":","`","TracingState","`","with","an","empty","`","TensorArray","`","in","its","`","trace_state","`","field","."]
code: def initialize(
      self, initial_chain_state, initial_kernel_results=None
  ):
      """Initializes a `TracingState` using previously defined metadata.
  
      Both the `initial_chain_state` and `initial_kernel_results` do not count
      as a sample, and hence, will not be traced. This is a deliberate decision
      that ensures consistency across sampling procedures.
  
      Args:
        initial_chain_state: A (possibly nested) structure of `Tensor`s or Python
          `list`s of `Tensor`s representing the current state(s) of the Markov
          chain(s). It is used to infer the structure of future trace results.
        initial_kernel_results: A (possibly nested) structure of `Tensor`s
          representing internal calculations made in a related `TransitionKernel`.
          It is used to infer the structure of future trace results.
  
      Returns:
        state: `TracingState` with an empty `TensorArray` in its `trace_state`
          field.
      """
      with tf.name_scope(
          mcmc_util.make_name(
              self.name, "tracing_reducer", "initialize"
          )
      ):
          initial_chain_state = tf.nest.map_structure(
              tf.convert_to_tensor, initial_chain_state
          )
          trace_result = self.trace_fn(
              initial_chain_state, initial_kernel_results
          )
          if self.size is None:
              size = 0
              dynamic_size = True
          else:
              size = self.size
              dynamic_size = False
  
          def _map_body(trace_state):
              if not tf.is_tensor(trace_state):
                  trace_state = tf.convert_to_tensor(trace_state)
              return loop_util.tensor_array_from_element(
                  trace_state,
                  size=size,
                  dynamic_size=dynamic_size,
                  clear_after_read=False,
              )
  
          trace_state = tf.nest.map_structure(_map_body, trace_result)
          return TracingState(0, 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 2 for calling the TracingReducer obj's underlying member method code with example usage: obj.initialize(initial_chain_state, initial_kernel_results) and returns: TracingState, loop_util