code: def test_assign_moving_mean_variance(self):
      shape = [3, 2]
      true_mean = np.array([[0.0, 3.0]])
      true_stddev = np.array([[1.1, 0.5]])
  
      mean_var = tf.Variable(tf.zeros_like(true_mean))
      variance_var = tf.Variable(tf.zeros_like(true_stddev))
      zero_debias_count = tf.Variable(0)
      self.evaluate(
          [
              mean_var.initializer,
              variance_var.initializer,
              zero_debias_count.initializer,
          ]
      )
  
      def body(it, seed):
          next_seed, current_seed = samplers.split_seed(seed)
          x = tf.random.stateless_normal(
              shape, dtype=np.float64, seed=current_seed
          )
          x = true_stddev * x + true_mean
          ema, emv = moving_stats.assign_moving_mean_variance(
              x,
              mean_var,
              variance_var,
              zero_debias_count,
              decay=0.999,
              axis=-2,
          )
          with tf.control_dependencies([ema, emv]):
              return [it + 1, next_seed]
  
      # Run 10000 updates; moving averages should be near the true values.
      run_op, _ = tf.while_loop(
          cond=lambda it, _: it < 10000,
          body=body,
          loop_vars=[0, test_util.test_seed(sampler_type="stateless")],
          parallel_iterations=1,
      )
  
      with tf.control_dependencies([run_op]):
          [
              final_unbiased_mean,
              final_ubiased_variance,
          ] = moving_stats.moving_mean_variance_zero_debiased(
              mean_var, variance_var, zero_debias_count, decay=0.999
          )
          final_biased_mean = tf.convert_to_tensor(mean_var)
          final_biased_variance = tf.convert_to_tensor(variance_var)
          read_zero_debias_count = tf.convert_to_tensor(
              zero_debias_count
          )
  
      [
          final_biased_mean_,
          final_biased_variance_,
          final_unbiased_mean_,
          final_unbiased_variance_,
          read_zero_debias_count_,
      ] = self.evaluate(
          [
              final_biased_mean,
              final_biased_variance,
              final_unbiased_mean,
              final_ubiased_variance,
              read_zero_debias_count,
          ]
      )
  
      # Test that values are as expected.
      # What tolerance should we set for the mean?  The expential moving average
      # of iid Gaussians is itself a Gaussian with the same mean.  If you assume
      # the average has converged (i.e., forgotten its initialization), then you
      # can work out that the variance is multiplied by a factor of
      # (1-decay)/(1+decay); that is, the effective sample size is
      # (1+decay)/(1-decay).  To assure low flakiness, set the tolerance
      # to > 4 * stddev.  We compute:
      # `4 * stddev = 4 * true_stddev * sqrt((1-decay)/(1+decay))
      #             = 4 * true_stddev * sqrt(0.001/1.999)
      #             = 4 * [1.1, 0.5]  * 0.0223
      #             < 0.1`
      # so `atol=0.15` ensures failures are highly unlikely.
      self.assertAllClose(true_mean, final_biased_mean_, atol=0.15)
      # The tolerance for the variance is more annoying to derive.  The handwave
      # justiying `atol=0.15` goes like this: The variance is the mean of the
      # square; so one would think the variance of the estimator of the variance
      # is around the square of the variance.  Since the variance is around 1,
      # that square doesn't matter, and the same tolerance should do.  </handwave>
      self.assertAllClose(
          true_stddev**2.0, final_biased_variance_, atol=0.15
      )
  
      self.assertEqual(10000, read_zero_debias_count_)
      self.assertAllClose(
          final_biased_mean_, final_unbiased_mean_, atol=0.01, rtol=0.01
      )
      self.assertAllClose(
          final_biased_variance_,
          final_unbiased_variance_,
          atol=0.01,
          rtol=0.01,
      )
  
      # Change the mean, var then update some more. Moving averages should
      # re-converge.
      self.evaluate(
          [
              mean_var.assign(np.array([[-1.0, 2.0]])),
              variance_var.assign(np.array([[2.0, 1.0]])),
          ]
      )
      # Run 10000 updates; moving averages should be near the true values.
      run_op, _ = tf.while_loop(
          cond=lambda it, _: it < 10000,
          body=body,
          loop_vars=[0, test_util.test_seed(sampler_type="stateless")],
          parallel_iterations=1,
      )
      self.evaluate(run_op)
  
      [final_biased_mean_, final_biased_variance_] = self.evaluate(
          [mean_var, variance_var]
      )
      # Test that values are as expected.
      self.assertAllClose(true_mean, final_biased_mean_, atol=0.15)
      self.assertAllClose(
          true_stddev**2.0, final_biased_variance_, atol=0.15
      )
  
file_path: probability/tensorflow_probability/python/stats/moving_stats_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.stats import moving_stats
label_desc: Use image node_id 1 for calling the MovingReduceMeanVarianceTest obj's underlying member method code with example usage: obj.test_assign_moving_mean_variance() and returns: list