code: def pop(params: LinearParams, old: PopArtState, new: PopArtState):
      """Preserves outputs precisely.
  
      Args:
        params: The parameters of the linear to preserve.
        old: The old PopArt state.
        new: The new PopArt state.
  
      Returns:
        new parameters.
      """
      w_new = params["w"] * jnp.broadcast_to(
          old.scale / new.scale, params["w"].shape
      )
      b_new = (
          old.scale * params["b"] + old.shift - new.shift
      ) / new.scale
      params_new = dict(w=w_new, b=b_new)
      return params_new
  
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 6 for calling a global function with example usage: pop(params, old, new) and returns: params_new