code: def test_learn_scale_shift(self):
      num_outputs = 2
      initial_state, update = pop_art.popart(
          num_outputs, step_size=1e-1, scale_lb=1e-6, scale_ub=1e6
      )
      state = initial_state()
      params = get_constant_linear_params(num_outputs)
      targets = np.arange(6) - 3
      indices = np.asarray([0, 0, 0, 1, 1, 1])
      # Learn the parameters.
      for _ in range(10):
          _, state = update(params, state, targets, indices)
  
      expected_scale = np.std(targets[:3])
      expected_scale = np.asarray([expected_scale, expected_scale])
      expected_shift = np.asarray([-2.0, 1.0])
      # Loose tolerances; just get close.
      np.testing.assert_allclose(
          state.scale, expected_scale, atol=1e-1, rtol=1e-1
      )
      np.testing.assert_allclose(
          state.shift, expected_shift, atol=1e-1, rtol=1e-1
      )
  
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 4 for calling the PopArtTest obj's underlying member method code with example usage: obj.test_learn_scale_shift() without return types