code: def popart(
      num_outputs: int,
      step_size: float,
      scale_lb: float,
      scale_ub: float,
      axis_name=None,
  ):
      """Generates functions giving initial PopArt state and update rule.
  
      Args:
        num_outputs: The number of outputs generated by the linear we're preserving.
        step_size: The step size for learning the scale & shift parameters.
        scale_lb: Lower bound for the scale.
        scale_ub: Upper bound for the scale.
        axis_name: What axis to aggregate over, if str. If passed an iterable,
          aggregates over multiple axes. Defaults to no aggregation, i.e. None.
  
      Returns:
        A tuple of:
          initial_state: A function returning the initial PopArt state.
          popart_update: A function updating the PopArt state and parameters
          of the preceding linear.
      """
  
      def initial_state():
          return PopArtState(
              jnp.zeros([num_outputs]),
              jnp.ones([num_outputs]),
              jnp.ones([num_outputs]),
          )
  
      def popart_update(
          params: LinearParams,
          state: PopArtState,
          targets: Array,
          indices: Array,
      ) -> Tuple[LinearParams, PopArtState]:
          """Computes the PopArt update.
  
          Args:
            params: The parameters of the linear to preserve.
            state: The current PopArt state.
            targets: Values whose distribution to learn.
            indices: For each target, which shift and scale element to adjust.
  
          Returns:
            A tuple of:
              new_params: The new parameters of the linear, preserving outputs.
              new_state: The new PopArt state.
          """
          # Disables Popart if step_size is None
          if step_size is None:
              return params, state
  
          # Adaptively rescale targets.
          state_new = art(
              state,
              targets,
              indices,
              step_size,
              scale_lb,
              scale_ub,
              axis_name,
          )
          # Preserve outputs precisely.
          params_new = pop(params, state, state_new)
          return params_new, state_new
  
      return initial_state, popart_update
  
file_path: rlax/rlax/_src/pop_art.py
imports: import collections
  from typing import Mapping, Tuple
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 7 for calling a global function with example usage: popart(num_outputs, step_size, scale_lb, scale_ub, axis_name) and returns: initial_state,  popart_update, PopArtState, params_new,  state_new, params,  state