tfp.substrates.jax.stats.moving_mean_variance_zero_debiased

Compute zero debiased versions of moving_mean and moving_variance.

tfp.substrates.jax.stats.moving_mean_variance_zero_debiased(
    moving_mean,
    moving_variance=None,
    zero_debias_count=None,
    decay=0.99,
    name=None
)

Since moving_* variables initialized with 0s will be biased (toward 0), this function rescales the moving_mean and moving_variance by the factor 1 - decay**zero_debias_count, i.e., such that the moving_mean is unbiased. For more details, see [Kingma (2014)][1].

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.Variables moving_mean and moving_variance always store the unbiased calculation, regardless of setting this argument. To obtain unbiased calculations from these tf.Variables, 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., 'moving_mean_variance_zero_debiased').

moving_mean The zero debiased exponentially weighted moving mean.
moving_variance The zero debiased exponentially weighted moving variance.

TypeError if moving_mean does not have float type dtype.
TypeError if moving_mean, moving_variance, decay have different base_dtype.

References

[1]: Diederik P. Kingma, Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv preprint arXiv:1412.6980, 2014. https://arxiv.org/abs/1412.6980