class_docstr: A running covariance computation.
  
  The running covariance computation supports batching. The `event_ndims`
  parameter indicates the number of trailing dimensions to treat as part of
  the event, and to compute covariance across. The leading dimensions, if
  any, are treated as batch shape, and no cross terms are computed.
  
  For example, if the incoming samples have shape `[5, 7]`, the `event_ndims`
  selects among three different covariance computations:
  - `event_ndims=0` treats the samples as a `[5, 7]` batch of scalar random
    variables, and computes their variances in batch.  The shape of the result
    is `[5, 7]`.
  - `event_ndims=1` treats the samples as a `[5]` batch of vector random
    variables of shape `[7]`, and computes their covariances in batch.  The
    shape of the result is `[5, 7, 7]`.
  - `event_ndims=2` treats the samples as a single random variable of
    shape `[5, 7]` and computes its covariance.  The shape of the result
    is `[5, 7, 5, 7]`.
  
  `RunningCovariance` is meant to serve general streaming covariance needs.
  For a specialized version that fits streaming over MCMC samples, see
  `CovarianceReducer` in `tfp.experimental.mcmc`.
docstr: Update the `RunningCovariance` with a new sample.
  
  The update formula is from Philippe Pebay (2008) [1]. This implementation
  supports both batched and chunked covariance computation. A "batch" is the
  usual parallel computation, namely a batch of size N implies N independent
  covariance computations, each stepping one sample (or chunk) at a time. A
  "chunk" of size M implies incorporating M samples into a single covariance
  computation at once, which is more efficient than one by one.
  
  To further illustrate the difference between batching and chunking, consider
  the following example:
  
  ```python
  # treat as 3 samples from each of 5 independent vector random variables of
  # shape (2,)
  sample = tf.ones((3, 5, 2))
  running_cov = tfp.experimental.stats.RunningCovariance.from_shape(
      (5, 2), event_ndims=1)
  running_cov = running_cov.update(sample, axis=0)
  final_cov = running_cov.covariance()
  final_cov.shape # (5, 2, 2)
  ```
  
  Args:
    new_sample: Incoming sample with shape and dtype compatible with those
      used to form this `RunningCovariance`.
    axis: If chunking is desired, this is an integer that specifies the axis
      with chunked samples. For individual samples, set this to `None`. By
      default, samples are not chunked (`axis` is None).
  
  Returns:
    cov: Newly allocated `RunningCovariance` updated to include `new_sample`.
  
  #### References
  [1]: Philippe Pebay. Formulas for Robust, One-Pass Parallel Computation of
       Covariances and Arbitrary-Order Statistical Moments. _Technical Report
       SAND2008-6212_, 2008.
       https://prod-ng.sandia.gov/techlib-noauth/access-control.cgi/2008/086212.pdf
docstr_tok: ["Update","the","`","RunningCovariance","`","with","a","new","sample",".","The","update","formula","is","from","Philippe","Pebay","(","2008",")","[","1","]",".","This","implementation","supports","both","batched","and","chunked","covariance","computation",".","A","``","batch","''","is","the","usual","parallel","computation",",","namely","a","batch","of","size","N","implies","N","independent","covariance","computations",",","each","stepping","one","sample","(","or","chunk",")","at","a","time",".","A","''","chunk","''","of","size","M","implies","incorporating","M","samples","into","a","single","covariance","computation","at","once",",","which","is","more","efficient","than","one","by","one",".","To","further","illustrate","the","difference","between","batching","and","chunking",",","consider","the","following","example",":","``","`","python","#","treat","as","3","samples","from","each","of","5","independent","vector","random","variables","of","#","shape","(","2",",",")","sample","=","tf.ones","(","(","3",",","5",",","2",")",")","running_cov","=","tfp.experimental.stats.RunningCovariance.from_shape","(","(","5",",","2",")",",","event_ndims=1",")","running_cov","=","running_cov.update","(","sample",",","axis=0",")","final_cov","=","running_cov.covariance","(",")","final_cov.shape","#","(","5",",","2",",","2",")","``","`","Args",":","new_sample",":","Incoming","sample","with","shape","and","dtype","compatible","with","those","used","to","form","this","`","RunningCovariance","`",".","axis",":","If","chunking","is","desired",",","this","is","an","integer","that","specifies","the","axis","with","chunked","samples",".","For","individual","samples",",","set","this","to","`","None","`",".","By","default",",","samples","are","not","chunked","(","`","axis","`","is","None",")",".","Returns",":","cov",":","Newly","allocated","`","RunningCovariance","`","updated","to","include","`","new_sample","`",".","#","#","#","#","References","[","1","]",":","Philippe","Pebay",".","Formulas","for","Robust",",","One-Pass","Parallel","Computation","of","Covariances","and","Arbitrary-Order","Statistical","Moments",".","_Technical","Report","SAND2008-6212_",",","2008.","https",":","\/\/prod-ng.sandia.gov\/techlib-noauth\/access-control.cgi\/2008\/086212.pdf"]
code: def update(self, new_sample, axis=None):
      """Update the `RunningCovariance` with a new sample.
  
      The update formula is from Philippe Pebay (2008) [1]. This implementation
      supports both batched and chunked covariance computation. A "batch" is the
      usual parallel computation, namely a batch of size N implies N independent
      covariance computations, each stepping one sample (or chunk) at a time. A
      "chunk" of size M implies incorporating M samples into a single covariance
      computation at once, which is more efficient than one by one.
  
      To further illustrate the difference between batching and chunking, consider
      the following example:
  
      ```python
      # treat as 3 samples from each of 5 independent vector random variables of
      # shape (2,)
      sample = tf.ones((3, 5, 2))
      running_cov = tfp.experimental.stats.RunningCovariance.from_shape(
          (5, 2), event_ndims=1)
      running_cov = running_cov.update(sample, axis=0)
      final_cov = running_cov.covariance()
      final_cov.shape # (5, 2, 2)
      ```
  
      Args:
        new_sample: Incoming sample with shape and dtype compatible with those
          used to form this `RunningCovariance`.
        axis: If chunking is desired, this is an integer that specifies the axis
          with chunked samples. For individual samples, set this to `None`. By
          default, samples are not chunked (`axis` is None).
  
      Returns:
        cov: Newly allocated `RunningCovariance` updated to include `new_sample`.
  
      #### References
      [1]: Philippe Pebay. Formulas for Robust, One-Pass Parallel Computation of
           Covariances and Arbitrary-Order Statistical Moments. _Technical Report
           SAND2008-6212_, 2008.
           https://prod-ng.sandia.gov/techlib-noauth/access-control.cgi/2008/086212.pdf
      """
      with tf.name_scope(
          (self.name or "RunningCovariance") + ".update"
      ):
          dtype = self.mean.dtype
          new_sample = tf.cast(new_sample, dtype=dtype)
          if axis is not None:
              chunk_n = tf.cast(ps.shape(new_sample)[axis], dtype=dtype)
              chunk_mean = tf.math.reduce_mean(new_sample, axis=axis)
              chunk_delta_mean = new_sample - tf.expand_dims(
                  chunk_mean, axis=axis
              )
              chunk_sum_squared_residuals = tf.reduce_sum(
                  _batch_outer_product(
                      chunk_delta_mean, self.event_ndims
                  ),
                  axis=axis,
              )
          else:
              chunk_n = tf.ones((), dtype=dtype)
              chunk_mean = new_sample
              chunk_sum_squared_residuals = tf.zeros(
                  ps.shape(self.sum_squared_residuals), dtype=dtype
              )
  
          new_n = self.num_samples + chunk_n
          delta_mean = chunk_mean - self.mean
          new_mean = self.mean + chunk_n * delta_mean / new_n
          all_pairwise_deltas = _batch_outer_product(
              delta_mean, self.event_ndims
          )
          adj_factor = (
              self.num_samples * chunk_n / (self.num_samples + chunk_n)
          )
          new_sum_squared_residuals = (
              self.sum_squared_residuals
              + chunk_sum_squared_residuals
              + adj_factor * all_pairwise_deltas
          )
          return type(self)(
              new_n,
              new_mean,
              new_sum_squared_residuals,
              self.event_ndims,
          )
  
imports: import functools
  import math
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow.python.util import nest
label_desc: Use image node_id 6 for calling the RunningCovariance obj's underlying member method code with example usage: obj.update(new_sample, axis) and returns: type