class_docstr: `Reducer` that computes a running expectation.
  
  `ExpectationsReducer` calculates expectations over some arbitrary structure
  of `transform_fn`s. A `transform_fn` is a function that accepts a Markov
  chain sample and kernel results, and outputs the relevant value for
  expectation calculation. In other words, if we denote a `transform_fn`
  as f(x,y), `ExpectationsReducer` approximates E[f(x,y)] for all provided
  functions. The finalized expectation will also identically mimic the
  structure of `transform_fn`.
  
  As with all reducers, ExpectationsReducer does not hold state information;
  rather, it stores supplied metadata. Intermediate calculations are held in
  a state object, which is returned via `initialize` and `one_step` method
  calls.
class_docstr_tok: ["`","Reducer","`","that","computes","a","running","expectation",".","`","ExpectationsReducer","`","calculates","expectations","over","some","arbitrary","structure","of","`","transform_fn","`","s",".","A","`","transform_fn","`","is","a","function","that","accepts","a","Markov","chain","sample","and","kernel","results",",","and","outputs","the","relevant","value","for","expectation","calculation",".","In","other","words",",","if","we","denote","a","`","transform_fn","`","as","f","(","x",",","y",")",",","`","ExpectationsReducer","`","approximates","E","[","f","(","x",",","y",")","]","for","all","provided","functions",".","The","finalized","expectation","will","also","identically","mimic","the","structure","of","`","transform_fn","`",".","As","with","all","reducers",",","ExpectationsReducer","does","not","hold","state","information",";","rather",",","it","stores","supplied","metadata",".","Intermediate","calculations","are","held","in","a","state","object",",","which","is","returned","via","`","initialize","`","and","`","one_step","`","method","calls","."]
docstr: Update the `current_reducer_state` with a new chain state.
  
  Chunking semantics are specified by the `axis` parameter. If chunking is
  enabled (axis is not `None`), all elements along the specified `axis` will
  be treated as separate samples. If a single scalar value is provided for a
  non-scalar sample structure, that value will be used for all elements in the
  structure. If not, an identical structure must be provided.
  
  Args:
    new_chain_state: A (possibly nested) structure of incoming chain state(s)
      with shape and dtype compatible with those used to initialize the
      `current_reducer_state`.
    current_reducer_state: `ExpectationsReducerState` representing the current
      reducer state.
    previous_kernel_results: A (possibly nested) structure of `Tensor`s
      representing internal calculations made in a related
      `TransitionKernel`.
    axis: If chunking is desired, this is a (possibly nested) structure of
      integers that specifies the axis with chunked samples. For individual
      samples, set this to `None`. By default, samples are not chunked
      (`axis` is None).
  
  Returns:
    new_reducer_state: `ExpectationsReducerState` with updated running
      statistics. It tracks a running total and the number of processed
      samples.
docstr_tok: ["Update","the","`","current_reducer_state","`","with","a","new","chain","state",".","Chunking","semantics","are","specified","by","the","`","axis","`","parameter",".","If","chunking","is","enabled","(","axis","is","not","`","None","`",")",",","all","elements","along","the","specified","`","axis","`","will","be","treated","as","separate","samples",".","If","a","single","scalar","value","is","provided","for","a","non-scalar","sample","structure",",","that","value","will","be","used","for","all","elements","in","the","structure",".","If","not",",","an","identical","structure","must","be","provided",".","Args",":","new_chain_state",":","A","(","possibly","nested",")","structure","of","incoming","chain","state","(","s",")","with","shape","and","dtype","compatible","with","those","used","to","initialize","the","`","current_reducer_state","`",".","current_reducer_state",":","`","ExpectationsReducerState","`","representing","the","current","reducer","state",".","previous_kernel_results",":","A","(","possibly","nested",")","structure","of","`","Tensor","`","s","representing","internal","calculations","made","in","a","related","`","TransitionKernel","`",".","axis",":","If","chunking","is","desired",",","this","is","a","(","possibly","nested",")","structure","of","integers","that","specifies","the","axis","with","chunked","samples",".","For","individual","samples",",","set","this","to","`","None","`",".","By","default",",","samples","are","not","chunked","(","`","axis","`","is","None",")",".","Returns",":","new_reducer_state",":","`","ExpectationsReducerState","`","with","updated","running","statistics",".","It","tracks","a","running","total","and","the","number","of","processed","samples","."]
code: def one_step(
      self,
      new_chain_state,
      current_reducer_state,
      previous_kernel_results=None,
      axis=None,
  ):
      """Update the `current_reducer_state` with a new chain state.
  
      Chunking semantics are specified by the `axis` parameter. If chunking is
      enabled (axis is not `None`), all elements along the specified `axis` will
      be treated as separate samples. If a single scalar value is provided for a
      non-scalar sample structure, that value will be used for all elements in the
      structure. If not, an identical structure must be provided.
  
      Args:
        new_chain_state: A (possibly nested) structure of incoming chain state(s)
          with shape and dtype compatible with those used to initialize the
          `current_reducer_state`.
        current_reducer_state: `ExpectationsReducerState` representing the current
          reducer state.
        previous_kernel_results: A (possibly nested) structure of `Tensor`s
          representing internal calculations made in a related
          `TransitionKernel`.
        axis: If chunking is desired, this is a (possibly nested) structure of
          integers that specifies the axis with chunked samples. For individual
          samples, set this to `None`. By default, samples are not chunked
          (`axis` is None).
  
      Returns:
        new_reducer_state: `ExpectationsReducerState` with updated running
          statistics. It tracks a running total and the number of processed
          samples.
      """
      with tf.name_scope(
          mcmc_util.make_name(
              self.name, "expectations_reducer", "one_step"
          )
      ):
          new_chain_state = tf.nest.map_structure(
              tf.convert_to_tensor, new_chain_state
          )
          if previous_kernel_results is not None:
              previous_kernel_results = tf.nest.map_structure(
                  tf.convert_to_tensor,
                  previous_kernel_results,
                  expand_composites=True,
              )
          fn_results = tf.nest.map_structure(
              lambda fn: fn(new_chain_state, previous_kernel_results),
              self.transform_fn,
          )
          if not nest.is_nested(axis):
              axis = nest_util.broadcast_structure(fn_results, axis)
  
          def update(fn_results, state, axis):
              return state.update(fn_results, axis=axis)
  
          return ExpectationsReducerState(
              nest.map_structure(
                  update,
                  fn_results,
                  current_reducer_state.expectation_state,
                  axis,
                  check_types=False,
              )
          )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/expectations_reducer.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import reducer
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling the ExpectationsReducer obj's underlying member method code with example usage: obj.one_step(new_chain_state, current_reducer_state, previous_kernel_results, axis) and returns: ExpectationsReducerState, state