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: Returns a function that acts as f called when variables are replaced by their averages.
  
  Args:
      f: function to be called on the stored averages.
  
  Returns:
      A function that returns the output of calling f with stored variables replaced by
      their moving averages.
docstr_tok: ["Returns","a","function","that","acts","as","f","called","when","variables","are","replaced","by","their","averages",".","Args",":","f",":","function","to","be","called","on","the","stored","averages",".","Returns",":","A","function","that","returns","the","output","of","calling","f","with","stored","variables","replaced","by","their","moving","averages","."]
code: def replace_vars(self, f: Callable):
      """Returns a function that acts as f called when variables are replaced by their averages.
  
      Args:
          f: function to be called on the stored averages.
  
      Returns:
          A function that returns the output of calling f with stored variables replaced by
          their moving averages.
      """
  
      def wrap(*args, **kwargs):
          refs, new_values = self.refs_and_values()
          original_values = refs.tensors()
          refs.assign(new_values)
          try:
              return f(*args, **kwargs)
          finally:
              refs.assign(original_values)
  
      return wrap
  
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 4 for calling the ExponentialMovingAverage obj's underlying member method code with example usage: obj.replace_vars(f) and returns: wrap, f