code: def assign_moving_mean_variance(
      value,
      moving_mean,
      moving_variance=None,
      zero_debias_count=None,
      decay=0.99,
      axis=(),
      name=None,
  ):
      """Compute one update to the exponentially weighted moving mean and variance.
  
      The `value` updated exponentially weighted moving `moving_mean` and
      `moving_variance` are conceptually given by the following recurrence
      relations ([Welford (1962)][1]):
  
      ```python
      new_mean = old_mean + (1 - decay) * (value - old_mean)
      new_var  = old_var  + (1 - decay) * (value - old_mean) * (value - new_mean)
      ```
  
      This function implements the above recurrences in a numerically stable manner
      and also uses the `assign_add` op to allow concurrent lockless updates to the
      supplied variables.
  
      For additional references see
      [John D. Cook's Blog](https://www.johndcook.com/blog/standard_deviation),
      whereas we use `1 - decay = 1 / k`, and
      [Finch (2009; Eq.  143)][2], whereas we use `1 - decay = alpha`.
  
      Since variables that are initialized to a `0` value will be `0` biased,
      providing `zero_debias_count` triggers scaling the `moving_mean` and
      `moving_variance` by the factor of `1 - decay ** (zero_debias_count + 1)`.
      For more details, see `tfp.stats.moving_mean_variance_zero_debiased`.
  
      Args:
        value: `float`-like `Tensor` representing one or more streaming
          observations. When `axis` is non-empty `value ` is reduced (by mean) for
          updated `moving_mean` and `moving-variance`. Presumed to have same shape
          as `moving_mean` and `moving_variance`.
        moving_mean: `float`-like `tf.Variable` representing the exponentially
          weighted moving mean. Same shape as `moving_variance` and `value`. This
          function presumes the `tf.Variable` was created with all zero initial
          value(s).
        moving_variance: `float`-like `tf.Variable` representing the exponentially
          weighted moving variance. Same shape as `moving_mean` and `value`.  This
          function presumes the `tf.Variable` was created with all zero initial
          value(s).
          Default value: `None` (i.e., no moving variance is computed).
        zero_debias_count: `int`-like `tf.Variable` representing the number of times
          this function has been called on streaming input (*not* the number of
          reduced values used in this functions computation). When not `None` (the
          default) the returned values for `moving_mean` and `moving_variance` are
          "zero debiased", i.e., corrected for their presumed all zeros
          intialization. Note: the `tf.Variable`s `moving_mean` and
          `moving_variance` *always* store the unbiased calculation, regardless of
          setting this argument. To obtain unbiased calculations from these
          `tf.Variable`s, see `tfp.stats.moving_mean_variance_zero_debiased`.
          Default value: `None` (i.e., no zero debiasing calculation is made).
        decay: A `float`-like `Tensor` representing the moving mean decay. Typically
          close to `1.`, e.g., `0.99`.
          Default value: `0.99`.
        axis: The dimensions to reduce. If `()` (the default) no dimensions are
          reduced. If `None` all dimensions are reduced. Must be in the range
          `[-rank(value), rank(value))`.
          Default value: `()` (i.e., no reduction is made).
        name: Python `str` prepended to op names created by this function.
          Default value: `None` (i.e., 'assign_moving_mean_variance').
  
      Returns:
        moving_mean: The `value`-updated exponentially weighted moving mean.
          Debiased if `zero_debias_count is not None`.
        moving_variance: The `value`-updated exponentially weighted moving variance.
          Debiased if `zero_debias_count is not None`.
  
      Raises:
        TypeError: if `moving_mean` does not have float type `dtype`.
        TypeError: if `moving_mean`, `moving_variance`, `value`, `decay` have
          different `base_dtype`.
  
      #### Examples
  
      ```python
      import tensorflow as tf
      import tensorflow_probability as tfp
      tfd = tfp.distributions
      d = tfd.MultivariateNormalTriL(
          loc=[-1., 1.],
          scale_tril=tf.linalg.cholesky([[0.75, 0.05],
                                         [0.05, 0.5]]))
      d.mean()
      # ==> [-1.,  1.]
      d.variance()
      # ==> [0.75, 0.5]
      moving_mean = tf.Variable(tf.zeros(2))
      moving_variance = tf.Variable(tf.zeros(2))
      zero_debias_count = tf.Variable(0)
      for _ in range(100):
        m, v = tfp.stats.assign_moving_mean_variance(
          value=d.sample(3),
          moving_mean=moving_mean,
          moving_variance=moving_variance,
          zero_debias_count=zero_debias_count,
          decay=0.99,
          axis=-2)
        print(m.numpy(), v.numpy())
      # ==> [-1.0334632  0.9545268] [0.8126194 0.5118788]
      # ==> [-1.0293456   0.96070296] [0.8115873  0.50947404]
      # ...
      # ==> [-1.025172  0.96351 ] [0.7142789  0.48570773]
  
      m1, v1 = tfp.stats.moving_mean_variance_zero_debiased(
        moving_mean,
        moving_variance,
        zero_debias_count,
        decay=0.99)
      print(m.numpy(), v.numpy())
      # ==> [-1.025172  0.96351 ] [0.7142789  0.48570773]
      assert(all(m == m1))
      assert(all(v == v1))
      ```
  
      #### References
  
      [1]  B. P. Welford. Note on a Method for Calculating Corrected Sums of
           Squares and Products. Technometrics, Vol. 4, No. 3 (Aug., 1962), p419-20.
           http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.302.7503&rep=rep1&type=pdf
           http://www.jstor.org/stable/1266577
  
      [2]: Tony Finch. Incremental calculation of weighted mean and variance.
           _Technical Report_, 2009.
           http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
      """
      with tf.name_scope(name or "assign_moving_mean_variance"):
          base_dtype = dtype_util.base_dtype(moving_mean.dtype)
          if not dtype_util.is_floating(base_dtype):
              raise TypeError(
                  "Argument `moving_mean` is not float type (saw {}).".format(
                      dtype_util.name(moving_mean.dtype)
                  )
              )
  
          value = tf.convert_to_tensor(
              value, dtype=base_dtype, name="value"
          )
          decay = tf.convert_to_tensor(
              decay, dtype=base_dtype, name="decay"
          )
          # Force a read of `moving_mean` as  we'll need it twice.
          old_mean = tf.convert_to_tensor(
              moving_mean, dtype=base_dtype, name="old_mean"
          )
  
          updated_mean = moving_mean.assign_add(
              (1.0 - decay)
              * (tf.reduce_mean(value, axis=axis) - old_mean)
          )
  
          if zero_debias_count is not None:
              t = tf.cast(zero_debias_count.assign_add(1), base_dtype)
              # Could have used:
              #   bias_correction = -tf.math.expm1(t * tf.math.log(decay))
              # however since we expect decay to be nearly 1, we don't expect this to
              # bear a significant improvement, yet would incur higher computational
              # cost.
              bias_correction = 1.0 - decay**t
              with tf.control_dependencies([updated_mean]):
                  updated_mean = updated_mean / bias_correction
  
          if moving_variance is None:
              return updated_mean
  
          if base_dtype != dtype_util.base_dtype(moving_variance.dtype):
              raise TypeError(
                  "Arguments `moving_mean` and `moving_variance` do not "
                  "have same base `dtype` (saw {}, {}).".format(
                      dtype_util.name(moving_mean.dtype),
                      dtype_util.name(moving_variance.dtype),
                  )
              )
  
          if zero_debias_count is not None:
              old_t = tf.where(
                  t > 1.0, t - 1.0, tf.constant(np.inf, base_dtype)
              )
              old_bias_correction = 1.0 - decay**old_t
              old_mean = old_mean / old_bias_correction
  
          mean_sq_diff = tf.reduce_mean(
              tf.math.squared_difference(value, old_mean), axis=axis
          )
          updated_variance = moving_variance.assign_add(
              (1.0 - decay) * (decay * mean_sq_diff - moving_variance)
          )
  
          if zero_debias_count is not None:
              with tf.control_dependencies([updated_variance]):
                  updated_variance = updated_variance / bias_correction
  
          return updated_mean, updated_variance
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 1 for calling a global function with example usage: assign_moving_mean_variance(value, moving_mean, moving_variance, zero_debias_count, decay, axis, name) and returns: updated_mean,  updated_variance, updated_mean