code: def get_fake_pop_art_state(num_outputs):
      """Returns a fake PopArtState."""
      shift = np.arange(num_outputs).astype(np.float32) + 1
      scale = shift
      second_moment = np.square(scale) + np.square(shift)
      return pop_art.PopArtState(shift, scale, second_moment)
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import haiku
  from haiku import data_structures
  from haiku import initializers
  import jax
  import jax.nn
  import jax.numpy
  import numpy
  from rlax._src import pop_art
label_desc: Use image node_id 3 for calling a global function with example usage: get_fake_pop_art_state(num_outputs) and returns: pop_art