Indeed, implementing batch normalization in Numpy is quite simple. Below is an alternative implementation, using 'X' to represent the batch of input features: import numpy as np def batch_norm ( X , eps = 1e-8 ) : if len ( X . shape ) not in ( 2 , 4 ) : raise ValueError ( "Input array should be 2D or 4D" ) if len ( X . shape ) == 2 : mean = np . mean ( X , axis = 0 ) std_dev = np . std ( X , axis = 0 ) else : mean = np . mean ( X , axis = ( 0 , 2 , 3 ) , keepdims = True ) std_dev = np . std ( X , axis = ( 0 , 2 , 3 ) , keepdims = True ) X_normalized = ( X - mean ) / ( std_dev + eps ) return X_normalized In this example, 'eps' is included to prevent division by zero errors. The function initially verifies whether the input array is either 2D or 4D, which are standard dimensions in neural network layers (2D for fully connected layers and 4D for convolutional layers). For a 4D array, the mean and standard deviation are computed across the width and height of the feature map, while preserving the batch and feature dimensions. The normalized output retains the same shape as the input data.