class_docstr: Maintains an exponential moving average.
  
  This uses the Adam debiasing procedure.
  See https://arxiv.org/pdf/1412.6980.pdf for details.
class_docstr_tok: ["Maintains","an","exponential","moving","average",".","This","uses","the","Adam","debiasing","procedure",".","See","https",":","\/\/arxiv.org\/pdf\/1412.6980.pdf","for","details","."]
docstr: Updates the EMA and returns the new value.
  
  Args:
    value: The array-like object for which you would like to perform an
      exponential decay on.
    update_stats: A Boolean, whether to update the internal state
      of this object to reflect the input value. When `update_stats` is False
      the internal stats will remain unchanged.
  
  Returns:
    The exponentially weighted average of the input value.
docstr_tok: ["Updates","the","EMA","and","returns","the","new","value",".","Args",":","value",":","The","array-like","object","for","which","you","would","like","to","perform","an","exponential","decay","on",".","update_stats",":","A","Boolean",",","whether","to","update","the","internal","state","of","this","object","to","reflect","the","input","value",".","When","`","update_stats","`","is","False","the","internal","stats","will","remain","unchanged",".","Returns",":","The","exponentially","weighted","average","of","the","input","value","."]
code: def __call__(
      self,
      value: Union[float, jax.Array],
      update_stats: bool = True,
  ) -> jax.Array:
      """Updates the EMA and returns the new value.
  
      Args:
        value: The array-like object for which you would like to perform an
          exponential decay on.
        update_stats: A Boolean, whether to update the internal state
          of this object to reflect the input value. When `update_stats` is False
          the internal stats will remain unchanged.
  
      Returns:
        The exponentially weighted average of the input value.
      """
      if not isinstance(value, jax.Array):
          value = jnp.asarray(value)
  
      counter = hk.get_state(
          "counter",
          (),
          jnp.int32,
          init=hk.initializers.Constant(-self.warmup_length),
      )
      counter = counter + 1
  
      decay = jax.lax.convert_element_type(self.decay, value.dtype)
      if self.warmup_length > 0:
          decay = jax.lax.select(counter <= 0, 0.0, decay)
  
      one = jnp.ones([], value.dtype)
      hidden = hk.get_state(
          "hidden", value.shape, value.dtype, init=jnp.zeros
      )
      hidden = hidden * decay + value * (one - decay)
  
      average = hidden
      if self.zero_debias:
          average /= one - jnp.power(decay, counter)
  
      if update_stats:
          hk.set_state("counter", counter)
          hk.set_state("hidden", hidden)
          hk.set_state("average", average)
  
      return average
  
imports: import re
  from typing import Optional, Union
  import warnings
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import initializers
  from haiku._src import module
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling the ExponentialMovingAverage obj's underlying member method code with example usage: obj.__call__(value, update_stats) and returns: average