the forward calculation of batch normalization and layer normalization

normalization.py

import numpy as np

def batch_normalization(x, gamma, beta, eps=2e-5):
  mean = x.mean(axis=0)
  var = x.var(axis=0)
  x_hat = (x - mean) / np.sqrt(var + eps)
  return gamma * x_hat + beta

def layer_normalization(x, gamma, beta, eps=2e-5):
  mean = x.mean(axis=1)[:, np.newaxis]
  var = x.var(axis=1)[:, np.newaxis]
  x_hat = (x - mean) / np.sqrt(var + eps)
  return gamma * x_hat + beta

if __name__ == '__main__':
  batch_size = 32
  layer_size = 128
  x = np.random.rand(batch_size, layer_size)
  gamma = np.ones((batch_size, layer_size))
  beta = np.zeros((batch_size, layer_size))
  x_bn = batch_normalization(x, gamma, beta)
  print(x_bn.mean(axis=0))
  print(x_bn.var(axis=0))
  x_ln = layer_normalization(x, gamma, beta)
  print(x_ln.mean(axis=1))
  print(x_ln.var(axis=1))