code: def testBatchNormShapeNCHW(self):
      key = random.PRNGKey(0)
      # Regression test for https://github.com/google/jax/issues/461
      init_fun, apply_fun = stax.BatchNorm(axis=(0, 2, 3))
      input_shape = (4, 5, 6, 7)
  
      out_shape, params = init_fun(key, input_shape)
      inputs = random_inputs(self.rng(), input_shape).astype(
          params[0].dtype
      )
      out = apply_fun(params, inputs)
  
      self.assertEqual(out_shape, input_shape)
      beta, gamma = params
      self.assertEqual(beta.shape, (5,))
      self.assertEqual(gamma.shape, (5,))
      self.assertEqual(out_shape, out.shape)
  
imports: from absl.testing import absltest
  import numpy
  from jax._src import test_util
  from jax import random
  from jax.example_libraries import stax
  from jax import dtypes
  from jax import config
label_desc: Use image node_id 17 for calling the StaxTest obj's underlying member method code with example usage: obj.testBatchNormShapeNCHW() without return types