code: def batch_normalization(
      x, mean, variance, axis, offset=None, scale=None, epsilon=1e-3
  ):
      shape = [1] * len(x.shape)
      shape[axis] = mean.shape[0]
      mean = np.reshape(mean, shape)
      variance = np.reshape(variance, shape)
  
      inv = 1.0 / np.sqrt(variance + epsilon)
      if scale is not None:
          scale = np.reshape(scale, shape)
          inv = inv * scale
  
      res = -mean * inv
      if offset is not None:
          offset = np.reshape(offset, shape)
          res = res + offset
  
      return x * inv + res
  
imports: import jax
  import numpy
  from jax import lax
  from jax import numpy
  from keras.backend import standardize_data_format
  from keras.backend import standardize_dtype
  from keras.backend.common.backend_utils import compute_conv_transpose_padding_args_for_jax
  from keras.backend.config import epsilon
  from keras.backend.numpy.core import cast
  from keras.backend.numpy.core import convert_to_tensor
  from keras.backend.numpy.core import is_tensor
  from keras.utils.module_utils import scipy
label_desc: Use image node_id 32 for calling a global function with example usage: batch_normalization(x, mean, variance, axis, offset, scale, epsilon) and returns: unknown