code: def test_forward(self, momentum):
      layer = tl.BatchNorm(momentum=momentum)
      x = np.array(
          [
              [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
              [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]],
          ]
      ).astype(np.float32)
      _, _ = layer.init(shapes.signature(x))
      y = layer(x)
      running_mean, running_var, n_batches = layer.state
  
      fraction_old = momentum
      fraction_new = 1.0 - momentum
      mean_of_x = 11.5  # mean of range(24)
      var_of_x = 47.9167  # variance of range(24)
      np.testing.assert_allclose(
          running_mean, 0.0 * fraction_old + mean_of_x * fraction_new
      )
      np.testing.assert_allclose(
          running_var,
          1.0 * fraction_old + var_of_x * fraction_new,
          rtol=1e-6,
      )
      self.assertEqual(n_batches, 1)
      eps = 1e-5
      np.testing.assert_allclose(
          y, (x - mean_of_x) / np.sqrt(var_of_x + eps), rtol=1e-6
      )
  
file_path: trax/trax/layers/normalization_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  from trax import fastmath
  from trax import shapes
  import trax.layers
label_desc: Use image node_id 3 for calling the BatchNormTest obj's underlying member method code with example usage: obj.test_forward(momentum) without return types