code: def test_simple_training_nchw(self):
      layer = batch_norm.BatchNorm(
          create_scale=False,
          create_offset=False,
          decay_rate=0.9,
          data_format="NCHW",
      )
  
      inputs = np.ones([2, 5, 3, 3])
      scale = np.full((5, 1, 1), 0.5)
      offset = np.full((5, 1, 1), 2.0)
  
      result = layer(inputs, True, scale=scale, offset=offset)
      np.testing.assert_equal(result, np.full(inputs.shape, 2.0))
  
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 3 for calling the BatchNormTest obj's underlying member method code with example usage: obj.test_simple_training_nchw() without return types