code: def setUp(self):
      self.axis = 3
      self.inp = (
          np.arange(16)
          .reshape(2, 2, 2, 2)
          .transpose(0, 2, 3, 1)
          .astype("float32")
      )
      self.keras_model = keras.models.Sequential()
      self.epsilon = 10 ** (-3)
      self.gamma = np.array([2.0, 3.0]).astype("float32")
      self.beta = np.array([4.0, 5.0]).astype("float32")
      self.mean = np.array([3.0, 3.0]).astype("float32")
      self.var = np.array([4.0, 9.0]).astype("float32")
      batch_norm_layer = keras.layers.normalization.BatchNormalization(
          axis=self.axis, input_shape=(2, 2, 2)
      )
      self.keras_model.add(batch_norm_layer)
      batch_norm_layer.set_weights(
          np.array(
              [
                  self.gamma,  # gamma (scaling)
                  self.beta,  # beta (shift)
                  self.mean,  # mean
                  self.var,
              ]
          )
      )  # std
      self.keras_model.add(keras.layers.Flatten())
      dense_layer = keras.layers.Dense(output_dim=1)
      self.keras_model.add(dense_layer)
      dense_layer.set_weights(
          [
              np.ones((1, 8)).astype("float32").T,
              np.zeros(1).astype("float32"),
          ]
      )
      self.keras_model.compile(loss="mse", optimizer="sgd")
  
      keras_batchnorm_fprop_func = compile_func(
          [self.keras_model.layers[0].input, K.learning_phase()],
          self.keras_model.layers[0].output,
      )
      self.keras_batchnorm_fprop_func = compile_func(
          [self.keras_model.layers[0].input, K.learning_phase()],
          self.keras_model.layers[0].output,
      )
      self.keras_output_fprop_func = compile_func(
          [self.keras_model.layers[0].input, K.learning_phase()],
          self.keras_model.layers[-1].output,
      )
  
      grad = tf.gradients(
          tf.reduce_sum(self.keras_model.layers[-1].output[:, 0]),
          [self.keras_model.layers[0].input],
      )[0]
      self.grad_func = compile_func(
          [self.keras_model.layers[0].input, K.learning_phase()], grad
      )
  
      self.saved_file_path = "batchnorm_model.h5"
      if os.path.isfile(self.saved_file_path):
          os.remove(self.saved_file_path)
      self.keras_model.save(self.saved_file_path)
  
imports: from __future__ import division
  from __future__ import print_function
  from __future__ import absolute_import
  import unittest
  from unittest import skip
  import sys
  import os
  import numpy
  from deeplift.conversion import kerasapi_conversion
  import deeplift.layers
  from deeplift.layers import NonlinearMxtsMode, DenseMxtsMode
  from deeplift.util import compile_func
  import tensorflow
  import keras
  from keras import models
  from keras import backend
label_desc: Use image node_id 1 for calling the TestBatchNorm obj's underlying member method code with example usage: obj.setUp() without return types