class_docstr: Maintains exponential moving averages for each variable from provided VarCollection.
class_docstr_tok: ["Maintains","exponential","moving","averages","for","each","variable","from","provided","VarCollection","."]
docstr: Creates ExponentialMovingAverage instance with given hyperparameters.
  
  Args:
      momentum: the decay factor for the moving average.
      debias: bool indicating whether to use initialization bias correction.
      eps: small adjustment to prevent division by zero.
docstr_tok: ["Creates","ExponentialMovingAverage","instance","with","given","hyperparameters",".","Args",":","momentum",":","the","decay","factor","for","the","moving","average",".","debias",":","bool","indicating","whether","to","use","initialization","bias","correction",".","eps",":","small","adjustment","to","prevent","division","by","zero","."]
code: def __init__(
      self,
      vc: VarCollection,
      momentum: float = 0.999,
      debias: bool = False,
      eps: float = 1e-6,
  ):
      """Creates ExponentialMovingAverage instance with given hyperparameters.
  
      Args:
          momentum: the decay factor for the moving average.
          debias: bool indicating whether to use initialization bias correction.
          eps: small adjustment to prevent division by zero.
      """
      self.momentum = momentum
      self.debias = debias
      self.eps = eps
      self.step = StateVar(
          jn.array(0, jn.uint32), reduce=lambda x: x[0]
      )
      # Deduplicate variables and skip RandomState vars since they cannot be averaged.
      trainable, non_trainable = (
          {},
          {},
      )  # Use dicts since they are ordered since python >= 3.6
      for v in vc:
          if isinstance(v, RandomState):
              continue
          if isinstance(v, TrainRef):
              v = v.ref
          if isinstance(v, TrainVar):
              trainable[id(v)] = v
          else:
              non_trainable[id(v)] = v
      self.refs = ModuleList(
          list(non_trainable.values())
          + [TrainRef(v) for v in trainable.values()]
      )
      self.m = ModuleList(
          StateVar(jn.zeros_like(x.value)) for x in self.refs
      )
  
file_path: objax/objax/optimizer/ema.py
imports: from typing import Callable, Tuple, List
  import jax.numpy
  from objax.module import Module, ModuleList
  from objax.typing import JaxArray
  from objax.util import class_name
  from objax.variable import RandomState, TrainRef, StateVar, TrainVar, VarCollection
label_desc: Use image node_id 1 to create a new ExponentialMovingAverage object from inherited base classes: Module with example: obj = ExponentialMovingAverage(vc, momentum, debias, eps)