code: def create_ema(decay=0.999, pmean_axis_name=None):
      """An updater of moments.
  
      Given a `tree` it will track first and second moments of the leaves.
  
      Args:
        decay: The decay of the moments. I.e., the learning rate is `1 - decay`.
        pmean_axis_name: If not None, use lax.pmean to average the moment updates.
  
      Returns:
        Two functions: `(init_state, update_moments)`.
      """
  
      def init_state(template_tree):
          zeros = jax.tree_map(
              lambda x: jnp.zeros_like(jnp.mean(x)), template_tree
          )
          scalar_zero = jnp.ones([], dtype=jnp.float32)
          return EmaState(mu=zeros, nu=zeros, decay_product=scalar_zero)
  
      def _update(moment, value):
          mean = jnp.mean(value)
          # Compute the mean across all learner devices involved in the `pmap`.
          if pmean_axis_name is not None:
              mean = jax.lax.pmean(mean, axis_name=pmean_axis_name)
          return decay * moment + (1 - decay) * mean
  
      def update_moments(tree, state):
          squared_tree = jax.tree_map(jnp.square, tree)
          mu = jax.tree_map(_update, state.mu, tree)
          nu = jax.tree_map(_update, state.nu, squared_tree)
          state = EmaState(
              mu=mu, nu=nu, decay_product=state.decay_product * decay
          )
          return state.debiased_moments(), state
  
      return init_state, update_moments
  
file_path: rlax/rlax/_src/moving_averages.py
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: create_ema(decay, pmean_axis_name) and returns: init_state,  update_moments, EmaState, unknown, state,  state