Add functionality to calculate KL-divergence of the weight and bias distributions

[RafayAK]

Adding/logging KL-divergence metrics of the weights and bias distributions before and after training can show how the distribution changes with regards to training as the number of samples (both augmented and unaugmented change)