Implement batch normalization for training and inference, including batch statistics and running-statistic updates.