code: def test_nested_normalize(self):
      state = running_statistics.init_state(
          {
              "a": specs.Array((5,), jnp.float32),
              "b": specs.Array((2,), jnp.float32),
          }
      )
  
      x1 = {
          "a": jnp.arange(20, dtype=jnp.float32).reshape(2, 2, 5),
          "b": jnp.arange(8, dtype=jnp.float32).reshape(2, 2, 2),
      }
      x2 = {
          "a": jnp.arange(20, dtype=jnp.float32).reshape(2, 2, 5) + 20,
          "b": jnp.arange(8, dtype=jnp.float32).reshape(2, 2, 2) + 8,
      }
      x3 = {
          "a": jnp.arange(40, dtype=jnp.float32).reshape(4, 2, 5),
          "b": jnp.arange(16, dtype=jnp.float32).reshape(4, 2, 2),
      }
  
      state = update_and_validate(state, x1)
      state = update_and_validate(state, x2)
      state = update_and_validate(state, x3)
      normalized = running_statistics.normalize(x3, state)
  
      mean = tree.map_structure(
          lambda x: jnp.mean(x, axis=(0, 1)), normalized
      )
      std = tree.map_structure(
          lambda x: jnp.std(x, axis=(0, 1)), normalized
      )
      tree.map_structure(
          lambda x: self.assert_allclose(x, jnp.zeros_like(x)), mean
      )
      tree.map_structure(
          lambda x: self.assert_allclose(x, jnp.ones_like(x)), std
      )
  
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 7 for calling the RunningStatisticsTest obj's underlying member method code with example usage: obj.test_nested_normalize() without return types