class_docstr: `Reducer` that computes a running covariance.
  
  By default, `CovarianceReducer` computes covariance directly over the
  latent state, ignoring kernel results; however, it can also compute
  a running covariance on samples evaluated on some arbitrary structure of
  `transform_fn`s. A `trasform_fn` is defined as any function that accepts
  the chain state and kernel results, and ouptuts a `Tensor` or nested
  structure of `Tensor`s to compute the covariance over. To be explicit,
  if we denote a `transform_fn` as f(x, y), then `CovarianceReducer` will
  compute Cov(f(x, y)). If some structure of `transform_fn`s are given,
  the final statistic (as returned by the `finalize` method) will mimic
  the results of that structure.
  
  As with all reducers, CovarianceReducer 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.
  
  `CovarianceReducer` is meant to fit into the larger Streaming MCMC
  framework. `RunningCovariance` in `tfp.experimental.stats` is better suited
  for more generic streaming covariance needs.
  
  There are two ways to stream over MCMC samples. The first is to manually
  wrap `CovarianceReducer` in a `tfp.experimental.mcmc.WithReductions`
  `TransitionKernel`. Doing so enables a wide variety of possible compositions.
  For example, to perform covariance calculation on samples that survive
  thinning and burn-in, `WithReductions` should wrap around an appropriate
  `SampleDiscardingKernel`.
  
  If the sole objective is to estimate covariance, it may be more convenient to
  use a pre-defined driver like `tfp.experimental.mcmc.sample_fold`.
  `sample_fold` will automatically create the Transition Kernel onion, apply
  `kernel` samples to update reducer states, and `finalize` calculations.
  
  ```
  python
  
  kernel = ...
  reducer = tfp.experimental.mcmc.CovarianceReducer()
  covariance_estimate, _, _ = tfp.experimental.mcmc.sample_fold(
      num_steps=...,
      current_state=...,
      previous_kernel_results=...,
      kernel=kernel,
      reducers=reducer,
  )
  covariance_estimate  # estimate of covariance (not `CovarianceReducerState`)
  ```
docstr: Update the `current_reducer_state` with a new chain state.
  
  Chunking semantics are similar to those of batching and 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: `CovarianceReducerState`s representing the current
      state of the running covariance.
    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: `CovarianceReducerState` with updated running
      statistics. Its `cov_state` field has an identical structure to the
      results of `self.transform_fn`. Each of the individual values in that
      structure subsequently mimics the structure of `current_reducer_state`.
docstr_tok: ["Update","the","`","current_reducer_state","`","with","a","new","chain","state",".","Chunking","semantics","are","similar","to","those","of","batching","and","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",":","`","CovarianceReducerState","`","s","representing","the","current","state","of","the","running","covariance",".","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",":","`","CovarianceReducerState","`","with","updated","running","statistics",".","Its","`","cov_state","`","field","has","an","identical","structure","to","the","results","of","`","self.transform_fn","`",".","Each","of","the","individual","values","in","that","structure","subsequently","mimics","the","structure","of","`","current_reducer_state","`","."]
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 similar to those of batching and 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: `CovarianceReducerState`s representing the current
          state of the running covariance.
        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: `CovarianceReducerState` with updated running
          statistics. Its `cov_state` field has an identical structure to the
          results of `self.transform_fn`. Each of the individual values in that
          structure subsequently mimics the structure of `current_reducer_state`.
      """
      with tf.name_scope(
          mcmc_util.make_name(
              self.name, "covariance_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)
          running_covariances = tf.nest.map_structure(
              lambda cov, *args: cov.update(*args),
              current_reducer_state.cov_state,
              fn_results,
              axis,
              check_types=False,
          )
          return CovarianceReducerState(running_covariances)
  
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 CovarianceReducer obj's underlying member method code with example usage: obj.one_step(new_chain_state, current_reducer_state, previous_kernel_results, axis) and returns: CovarianceReducerState