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","."]
code: def __call__(self, tree, update_stats=True):
      def maybe_ema(k, v):
          if self.ignore_regex and re.match(self.ignore_regex, k):
              return v
          else:
              ema_name = k.replace("/", "__").replace("~", "_tilde_")
              return ExponentialMovingAverage(
                  self.decay,
                  self.zero_debias,
                  self.warmup_length,
                  name=ema_name,
              )(v, update_stats=update_stats)
  
      # We want to potentially replace params with EMA'd versions.
      new_values = {}
      for module_name, param_dict in tree.items():
          new_values[module_name] = {
              k: maybe_ema("/".join([module_name, k]), v)
              for k, v in param_dict.items()
          }
      return hk.data_structures.to_haiku_dict(new_values)
  
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 2 for calling the EMAParamsTree obj's underlying member method code with example usage: obj.__call__(tree, update_stats) and returns: hk, v, ExponentialMovingAverage