code: def test_correctness(
      self,
      axis,
      input_shape,
      moving_mean_initializer,
      moving_variance_initializer,
  ):
      # Training
      layer = layers.BatchNormalization(
          axis=axis,
          momentum=0,
          moving_mean_initializer=moving_mean_initializer,
          moving_variance_initializer=moving_variance_initializer,
      )
      # Random data centered on 5.0, variance 10.0
      x = np.random.normal(loc=5.0, scale=10.0, size=input_shape)
      out = x
      for _ in range(3):
          out = layer(out, training=True)
  
      # Assert the normalization is correct.
      broadcast_shape = [1] * len(input_shape)
      broadcast_shape[axis] = input_shape[axis]
      out = backend.convert_to_numpy(out)
      out -= np.reshape(
          backend.convert_to_numpy(layer.beta), broadcast_shape
      )
      out /= np.reshape(
          backend.convert_to_numpy(layer.gamma), broadcast_shape
      )
  
      reduction_axes = list(range(len(input_shape)))
      del reduction_axes[axis]
      reduction_axes = tuple(reduction_axes)
      self.assertAllClose(
          np.mean(out, axis=reduction_axes), 0.0, atol=1e-3
      )
      self.assertAllClose(
          np.std(out, axis=reduction_axes), 1.0, atol=1e-3
      )
      self.assertAllClose(layer.moving_mean, 0.0, atol=1e-3)
      self.assertAllClose(layer.moving_variance, 1.0, atol=1e-3)
  
      # Inference done before training shouldn't match.
      inference_out = layer(x, training=False)
      training_out = layer(x, training=True)
      self.assertNotAllClose(inference_out, training_out)
  
      # Since momentum is zero, inference after training should match.
      training_out = layer(x, training=True)
      inference_out = layer(x, training=False)
      self.assertAllClose(inference_out, training_out)
  
      # Masked result with no training should not differ
      x[:, 1, :] = 0.0
      unmasked_out = layer(x, training=False)
      masked = layers.Masking()(x)
      masked_out = layer(masked, training=False)
      self.assertAllClose(unmasked_out, masked_out)
  
      # Masked result should differ from unmasked result
      unmasked_out = layer(x, training=False)
      x[:, 1, :] = 0.0
      masked = layers.Masking()(x)
      masked_out = layer(masked, training=True)
      self.assertNotAllClose(unmasked_out, masked_out)
  
file_path: keras/keras/layers/normalization/batch_normalization_test.py
imports: import numpy
  import pytest
  from absl.testing import parameterized
  from keras import backend
  from keras import layers
  from keras import testing
  from keras.losses import MeanSquaredError
  from keras.models import Model
label_desc: Use image node_id 2 for calling the BatchNormalizationTest obj's underlying member method code with example usage: obj.test_correctness(axis, input_shape, moving_mean_initializer, moving_variance_initializer) without return types