code: def assign_log_moving_mean_exp(
      log_value,
      moving_log_mean_exp,
      zero_debias_count=None,
      decay=0.99,
      name=None,
  ):
      """Compute the log of the exponentially weighted moving mean of the exp.
  
      If `log_value` is a draw from a stationary random variable, this function
      approximates `log(E[exp(log_value)])`, i.e., a weighted log-sum-exp. More
      precisely, a `tf.Variable`, `moving_log_mean_exp`, is updated by `log_value`
      using the following identity:
  
      ```none
      moving_log_mean_exp =
      = log(decay exp(moving_log_mean_exp) + (1 - decay) exp(log_value))
      = log(exp(moving_log_mean_exp + log(decay)) + exp(log_value + log1p(-decay)))
      = moving_log_mean_exp
        + log(  exp(moving_log_mean_exp   - moving_log_mean_exp + log(decay))
              + exp(log_value - moving_log_mean_exp + log1p(-decay)))
      = moving_log_mean_exp
        + log_sum_exp([log(decay), log_value - moving_log_mean_exp +
        log1p(-decay)]).
      ```
  
      In addition to numerical stability, this formulation is advantageous because
      `moving_log_mean_exp` can be updated in a lock-free manner, i.e., using
      `assign_add`. (Note: the updates are not thread-safe; it's just that the
      update to the tf.Variable is presumed efficient due to being lock-free.)
  
      Args:
        log_value: `float`-like `Tensor` representing a new (streaming) observation.
          Same shape as `moving_log_mean_exp`.
        moving_log_mean_exp: `float`-like `Variable` representing the log of the
          exponentially weighted moving mean of the exp. Same shape as `log_value`.
        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`.
        name: Python `str` prepended to op names created by this function.
          Default value: `None` (i.e., 'assign_log_moving_mean_exp').
  
      Returns:
        moving_log_mean_exp: A reference to the input 'Variable' tensor with the
          `log_value`-updated log of the exponentially weighted moving mean of exp.
  
      Raises:
        TypeError: if `moving_log_mean_exp` does not have float type `dtype`.
        TypeError: if `moving_log_mean_exp`, `log_value`, `decay` have different
          `base_dtype`.
      """
      if zero_debias_count is not None:
          raise NotImplementedError(
              "Argument `zero_debias_count` is not yet supported. If you require "
              "this feature please create a new issue on "
              "`https://github.com/tensorflow/probability` or email "
              "`tfprobability@tensorflow.org`."
          )
      with tf.name_scope(name or "assign_log_moving_mean_exp"):
          # We want to update the variable in a numerically stable and lock-free way.
          # To do this, observe that variable `x` updated by `v` is:
          # x = log(w exp(x) + (1-w) exp(v))
          #   = log(exp(x + log(w)) + exp(v + log1p(-w)))
          #   = x + log(exp(x - x + log(w)) + exp(v - x + log1p(-w)))
          #   = x + lse([log(w), v - x + log1p(-w)])
          base_dtype = dtype_util.base_dtype(moving_log_mean_exp.dtype)
          if not dtype_util.is_floating(base_dtype):
              raise TypeError(
                  "Argument `moving_log_mean_exp` is not float type (saw {}).".format(
                      dtype_util.name(moving_log_mean_exp.dtype)
                  )
              )
          log_value = tf.convert_to_tensor(
              log_value, dtype=base_dtype, name="log_value"
          )
          decay = tf.convert_to_tensor(
              decay, dtype=base_dtype, name="decay"
          )
          delta = (log_value - moving_log_mean_exp)[tf.newaxis, ...]
          x = tf.concat(
              [
                  tf.broadcast_to(
                      tf.math.log(decay),
                      ps.broadcast_shape(
                          ps.shape(decay), ps.shape(delta)
                      ),
                  ),
                  delta + tf.math.log1p(-decay),
              ],
              axis=0,
          )
          update = tf.reduce_logsumexp(x, axis=0)
          return moving_log_mean_exp.assign_add(update)
  
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 3 for calling a global function with example usage: assign_log_moving_mean_exp(log_value, moving_log_mean_exp, zero_debias_count, decay, name) and returns: moving_log_mean_exp