code: def test_pmap_update_nested(self):
      local_device_count = jax.local_device_count()
      state = running_statistics.init_state(
          {
              "a": specs.Array((5,), jnp.float32),
              "b": specs.Array((2,), jnp.float32),
          }
      )
  
      x = {
          "a": (
              jnp.arange(15 * local_device_count, dtype=jnp.float32)
          ).reshape(local_device_count, 3, 5),
          "b": (
              jnp.arange(6 * local_device_count, dtype=jnp.float32)
          ).reshape(local_device_count, 3, 2),
      }
  
      devices = jax.local_devices()
      state = jax.device_put_replicated(state, devices)
      pmap_axis_name = "i"
      state = jax.pmap(
          functools.partial(
              update_and_validate, pmap_axis_name=pmap_axis_name
          ),
          pmap_axis_name,
      )(state, x)
      state = jax.pmap(
          functools.partial(
              update_and_validate, pmap_axis_name=pmap_axis_name
          ),
          pmap_axis_name,
      )(state, x)
      normalized = jax.pmap(running_statistics.normalize)(x, 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 10 for calling the RunningStatisticsTest obj's underlying member method code with example usage: obj.test_pmap_update_nested() without return types