code: def test_simple_training_cross_replica_axis(self):
      ldc = jax.local_device_count()
  
      def f(x, is_training=True):
          return batch_norm.BatchNorm(
              create_scale=False,
              create_offset=False,
              decay_rate=0.9,
              cross_replica_axis="i",
          )(x, is_training=is_training)
  
      f = transform.transform_with_state(f)
  
      inputs = np.arange(ldc * 4).reshape(ldc, 4)
      key = jax.random.PRNGKey(42)
      key = jnp.broadcast_to(key, (ldc, *key.shape))
      params, state = jax.pmap(f.init, axis_name="i")(key, inputs)
      result, _ = jax.pmap(f.apply, axis_name="i")(
          params, state, key, inputs
      )
  
      mean = np.mean(inputs, axis=0)
      std = np.std(inputs, axis=0) + 1e-10
      expected = (inputs - mean) / std
  
      np.testing.assert_array_almost_equal(result, expected)
  
imports: import os
  from absl.testing import absltest
  from haiku._src import base
  from haiku._src import batch_norm
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 5 for calling the BatchNormTest obj's underlying member method code with example usage: obj.test_simple_training_cross_replica_axis() and returns: batch_norm