code: def art(
      state: PopArtState,
      targets: Array,
      indices: Array,
      step_size: float,
      scale_lb: float,
      scale_ub: float,
      axis_name=None,
  ) -> PopArtState:
      """Adaptively rescale targets.
  
      Args:
        state: The PopArt summary stats.
        targets: targets which are rescaled.
        indices: Which indices of the state to use.
        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:
        New popart state which can be used to rescale targets.
      """
      assert targets.shape == indices.shape
      assert jnp.issubdtype(indices.dtype, jnp.integer)
  
      # Update shift.
      shift_gather = state.shift[indices]
      shift_update = step_size * (targets - shift_gather)
      shift_new = _cross_replica_scatter_add(
          state.shift, indices, shift_update, axis_name
      )
  
      # Update second moment.
      second_moment_gather = state.second_moment[indices]
      second_moment_update = step_size * (
          jnp.square(targets) - second_moment_gather
      )
      second_moment_new = _cross_replica_scatter_add(
          state.second_moment, indices, second_moment_update, axis_name
      )
  
      # Derive scale (stdev) from second moment and mean.
      scale_sq = second_moment_new - jnp.square(shift_new)
      scale_sq = jnp.clip(scale_sq, scale_lb**2, scale_ub**2)
      scale_new = jnp.sqrt(scale_sq)
  
      state_new = PopArtState(shift_new, scale_new, second_moment_new)
  
      # Prevent gradients propagating back through the state.
      state_new = jax.tree_map(jax.lax.stop_gradient, state_new)
  
      return state_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 5 for calling a global function with example usage: art(state, targets, indices, step_size, scale_lb, scale_ub, axis_name) and returns: state_new