code: def test_normalize_config(self):
      x = jnp.arange(200, dtype=jnp.float32).reshape(20, 2, 5)
      x_split = jnp.split(x, 5, axis=0)
  
      y = jnp.arange(160, dtype=jnp.float32).reshape(20, 2, 4)
      y_split = jnp.split(y, 5, axis=0)
  
      z = {"a": x, "b": y}
  
      z_split = [{"a": xx, "b": yy} for xx, yy in zip(x_split, y_split)]
  
      update = jax.jit(
          running_statistics.update, static_argnames=("config",)
      )
  
      config = running_statistics.NestStatisticsConfig((("a",),))
      state = running_statistics.init_state(
          {
              "a": specs.Array((5,), jnp.float32),
              "b": specs.Array((4,), jnp.float32),
          }
      )
      # Test initialization from the first element.
      state = update(state, z_split[0], config=config)
      state = update(state, z_split[1], config=config)
      state = update(state, z_split[2], config=config)
      state = update(state, z_split[3], config=config)
      state = update(state, z_split[4], config=config)
  
      normalize = jax.jit(running_statistics.normalize)
      normalized = normalize(z, state)
  
      for key in normalized:
          mean = jnp.mean(normalized[key], axis=(0, 1))
          std = jnp.std(normalized[key], axis=(0, 1))
          if key == "a":
              self.assert_allclose(
                  mean,
                  jnp.zeros_like(mean),
                  err_msg=f"key:{key} mean:{mean} normalized:{normalized[key]}",
              )
              self.assert_allclose(
                  std,
                  jnp.ones_like(std),
                  err_msg=f"key:{key} std:{std} normalized:{normalized[key]}",
              )
          else:
              assert key == "b"
              np.testing.assert_array_equal(
                  normalized[key],
                  z[key],
                  err_msg=f"z:{z[key]} normalized:{normalized[key]}",
              )
  
imports: import functools
  import math
  from typing import NamedTuple
  from acme import specs
  from acme.jax import running_statistics
  import jax
  from jax import config
  import jax.numpy
  import numpy
  import tree
  from absl.testing import absltest
label_desc: Use image node_id 13 for calling the RunningStatisticsTest obj's underlying member method code with example usage: obj.test_normalize_config() without return types