code: def test_weights(self):
      state = running_statistics.init_state(
          specs.Array((), jnp.float32)
      )
  
      x = jnp.arange(5, dtype=jnp.float32)
      x_weights = jnp.ones_like(x)
      y = 2 * x + 5
      y_weights = 2 * x_weights
      z = jnp.concatenate([x, y])
      weights = jnp.concatenate([x_weights, y_weights])
  
      state = update_and_validate(state, z, weights=weights)
  
      self.assertEqual(state.mean, (jnp.mean(x) + 2 * jnp.mean(y)) / 3)
      big_z = jnp.concatenate([x, y, y])
      normalized = running_statistics.normalize(big_z, state)
      self.assertAlmostEqual(jnp.mean(normalized), 0.0, places=6)
      self.assertAlmostEqual(jnp.std(normalized), 1.0, places=6)
  
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 12 for calling the RunningStatisticsTest obj's underlying member method code with example usage: obj.test_weights() without return types