Missing sum in Epistemic and Aleatoric computation

[Zimba96]

Hey there,

I recently tried to implement a Bayesian Net myself and oriented myself a little on your implementation. When coming across the computation of the epistemic and aleatoric, I noticed that your implementation in uncertainty_estimation.py seems to differ a little from the formula in your paper.
In the paper, it says:

However, I'm missing the sums in your implementation, e. g. when you calculate the epistemic and you divide by T but not summing over all t's (same with the second part of the aleatoric):

temp = p_hat - np.expand_dims(p_bar, 0)
epistemic = np.dot(temp.T, temp) / T
epistemic = np.diag(epistemic)

aleatoric = np.diag(p_bar) - (np.dot(p_hat.T, p_hat) / T)
aleatoric = np.diag(aleatoric)

Or does the following np.diag somehow replaces the summation?

There's another short question I'd like to ask:
This code outputs the epistemic and aleatoric as vectors of size [number_of_classes], right? Does that mean that the vector represents the epistemic and aleatoric for the predictions of each class? And when I want to obtain the epistemic and aleatoric of my resulting prediction (the class with the max. softmax/softplus-value), do I then only need to extract the aleatoric and epistemic-value corresponding to that class from the epistemic and aleatoric vectors?

Thanks in advance for your answers!