code: def test_outputs_preserved(self):
      num_outputs = 2
      initial_state, update = pop_art.popart(
          num_outputs, step_size=1e-3, scale_lb=1e-6, scale_ub=1e6
      )
      state = initial_state()
      key = jax.random.PRNGKey(428)
  
      def net(x):
          linear = hk.Linear(
              num_outputs,
              b_init=initializers.RandomUniform(),
              name="head",
          )
          return linear(x)
  
      init_fn, apply_fn = hk.without_apply_rng(hk.transform(net))
      key, subkey1, subkey2 = jax.random.split(key, 3)
      fixed_data = jax.random.uniform(subkey1, (4, 3))
      params = init_fn(subkey2, fixed_data)
      initial_result = apply_fn(params, fixed_data)
      indices = np.asarray([0, 1, 0, 1, 0, 1, 0, 1])
      # Repeatedly update state and verify that params still preserve outputs.
      for _ in range(30):
          key, subkey1, subkey2 = jax.random.split(key, 3)
          targets = jax.random.uniform(subkey1, (8,))
          linear_params, state = update(
              params["head"], state, targets, indices
          )
          params = data_structures.to_mutable_dict(params)
          params["head"] = linear_params
  
          # Apply updated linear transformation and unnormalize outputs.
          transform = apply_fn(params, fixed_data)
          out = jnp.broadcast_to(
              state.scale, transform.shape
          ) * transform + jnp.broadcast_to(state.shift, transform.shape)
          np.testing.assert_allclose(initial_result, out, atol=1e-2)
  
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 7 for calling the PopArtTest obj's underlying member method code with example usage: obj.test_outputs_preserved() and returns: linear