code: def test_simple_training_normalized_axes(self):
      layer = batch_norm.BatchNorm(
          create_scale=False,
          create_offset=False,
          decay_rate=0.9,
          axis=[0, 2, 3],
      )  # Not the second axis.
  
      # This differs only in the second axis.
      inputs = np.stack(
          [2.0 * np.ones([5, 3, 3]), np.ones([5, 3, 3])], 1
      )
  
      result = layer(inputs, True)
  
      # Despite not all values being identical, treating slices from the first
      # axis separately leads to a fully normalized = equal array.
      np.testing.assert_equal(result, np.zeros(inputs.shape))
  
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 4 for calling the BatchNormTest obj's underlying member method code with example usage: obj.test_simple_training_normalized_axes() without return types