class_docstr: data-class holding the exponential moving average state.
docstr: Returns debiased moments as in Adam.
docstr_tok: ["Returns","debiased","moments","as","in","Adam","."]
code: def debiased_moments(self):
      """Returns debiased moments as in Adam."""
      tiny = jnp.finfo(self.decay_product).tiny
      debias = 1.0 / jnp.maximum(1 - self.decay_product, tiny)
      mean = jax.tree_map(lambda m1: m1 * debias, self.mu)
      # This computation of the variance may lose some numerical precision, if
      # the mean is not approximately zero.
      variance = jax.tree_map(
          lambda m2, m: jnp.maximum(0.0, m2 * debias - jnp.square(m)),
          self.nu,
          mean,
      )
      return EmaMoments(mean=mean, variance=variance)
  
file_path: rlax/rlax/_src/moving_averages.py
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the EmaState obj's underlying member method code with example usage: obj.debiased_moments() and returns: EmaMoments