class_docstr: Maintains an exponential moving average for all parameters in a tree.
  
  While ExponentialMovingAverage is meant to be applied to single parameters
  within a function, this class is meant to be applied to the entire tree of
  parameters for a function.
  
  Given a set of parameters for some network:
  
  >>> network_fn = lambda x: hk.Linear(10)(x)
  >>> x = jnp.ones([1, 1])
  >>> params = hk.transform(network_fn).init(jax.random.PRNGKey(428), x)
  
  You might use the EMAParamsTree like follows:
  
  >>> ema_fn = hk.transform_with_state(lambda x: hk.EMAParamsTree(0.2)(x))
  >>> _, ema_state = ema_fn.init(None, params)
  >>> ema_params, ema_state = ema_fn.apply(None, ema_state, None, params)
  
  Here, we are transforming a Haiku function and constructing its parameters via
  an init_fn as normal, but are creating a second transformed function which
  expects a tree of parameters as input. This function is then called with
  the current parameters as input, which then returns an identical tree with
  every parameter replaced with its exponentially decayed average. This
  ema_params object can then be passed into the `network_fn` as usual, and will
  cause it to run with EMA weights.
class_docstr_tok: ["Maintains","an","exponential","moving","average","for","all","parameters","in","a","tree",".","While","ExponentialMovingAverage","is","meant","to","be","applied","to","single","parameters","within","a","function",",","this","class","is","meant","to","be","applied","to","the","entire","tree","of","parameters","for","a","function",".","Given","a","set","of","parameters","for","some","network",":",">",">",">","network_fn","=","lambda","x",":","hk.Linear","(","10",")","(","x",")",">",">",">","x","=","jnp.ones","(","[","1",",","1","]",")",">",">",">","params","=","hk.transform","(","network_fn",")",".init","(","jax.random.PRNGKey","(","428",")",",","x",")","You","might","use","the","EMAParamsTree","like","follows",":",">",">",">","ema_fn","=","hk.transform_with_state","(","lambda","x",":","hk.EMAParamsTree","(","0.2",")","(","x",")",")",">",">",">","_",",","ema_state","=","ema_fn.init","(","None",",","params",")",">",">",">","ema_params",",","ema_state","=","ema_fn.apply","(","None",",","ema_state",",","None",",","params",")","Here",",","we","are","transforming","a","Haiku","function","and","constructing","its","parameters","via","an","init_fn","as","normal",",","but","are","creating","a","second","transformed","function","which","expects","a","tree","of","parameters","as","input",".","This","function","is","then","called","with","the","current","parameters","as","input",",","which","then","returns","an","identical","tree","with","every","parameter","replaced","with","its","exponentially","decayed","average",".","This","ema_params","object","can","then","be","passed","into","the","`","network_fn","`","as","usual",",","and","will","cause","it","to","run","with","EMA","weights","."]
docstr: Initializes an EMAParamsTree module.
  
  Args:
    decay: The chosen decay. Must in ``[0, 1)``. Values close to ``1`` result
      in slow decay; values close to ``0`` result in fast decay.
    zero_debias: Whether to run with zero-debiasing.
    warmup_length: A positive integer, EMA has no effect until
      the internal counter has reached `warmup_length` at which point the
      initial value for the decaying average is initialized to the input value
      after `warmup_length` iterations.
    ignore_regex: A string. Any parameter in the tree whose name matches this
      regex will not have any moving average applied to it. The empty string
      means this module will EMA all parameters.
    name: The name of the module.
docstr_tok: ["Initializes","an","EMAParamsTree","module",".","Args",":","decay",":","The","chosen","decay",".","Must","in","``","[","0",",","1",")","``",".","Values","close","to","``","1","``","result","in","slow","decay",";","values","close","to","``","0","``","result","in","fast","decay",".","zero_debias",":","Whether","to","run","with","zero-debiasing",".","warmup_length",":","A","positive","integer",",","EMA","has","no","effect","until","the","internal","counter","has","reached","`","warmup_length","`","at","which","point","the","initial","value","for","the","decaying","average","is","initialized","to","the","input","value","after","`","warmup_length","`","iterations",".","ignore_regex",":","A","string",".","Any","parameter","in","the","tree","whose","name","matches","this","regex","will","not","have","any","moving","average","applied","to","it",".","The","empty","string","means","this","module","will","EMA","all","parameters",".","name",":","The","name","of","the","module","."]
code: def __init__(
      self,
      decay,
      zero_debias: bool = True,
      warmup_length: int = 0,
      ignore_regex: str = "",
      name: Optional[str] = None,
  ):
      """Initializes an EMAParamsTree module.
  
      Args:
        decay: The chosen decay. Must in ``[0, 1)``. Values close to ``1`` result
          in slow decay; values close to ``0`` result in fast decay.
        zero_debias: Whether to run with zero-debiasing.
        warmup_length: A positive integer, EMA has no effect until
          the internal counter has reached `warmup_length` at which point the
          initial value for the decaying average is initialized to the input value
          after `warmup_length` iterations.
        ignore_regex: A string. Any parameter in the tree whose name matches this
          regex will not have any moving average applied to it. The empty string
          means this module will EMA all parameters.
        name: The name of the module.
      """
      super().__init__(name=name)
      self.decay = decay
      self.zero_debias = zero_debias
      self.warmup_length = warmup_length
      self.ignore_regex = ignore_regex
  
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 1 to create a new EMAParamsTree object from inherited base classes: hk with example: obj = EMAParamsTree(decay, zero_debias, warmup_length, ignore_regex, name)