code: def test_basic(self):
      data = jnp.arange(2 * 3 * 4, dtype=jnp.float32).reshape([2, 3, 4])
  
      norm = batch_norm.BatchNorm(True, True, 0.9)
      result = norm(data, is_training=True)
      result_0_replicated = jnp.broadcast_to(
          result[:, :, :1], result.shape
      )
      # Input data is symmetrical variance per-channel.
      np.testing.assert_allclose(result, result_0_replicated)
      # Running through again in test mode produces same output.
      np.testing.assert_allclose(
          norm(data, is_training=False), result, rtol=2e-2
      )
  
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 1 for calling the BatchNormTest obj's underlying member method code with example usage: obj.test_basic() without return types