How to implement an equation with "diag" and circled-times symbol (⊗) in Python

Question

The above equation is found in this paper (https://openreview.net/pdf?id=Sk_P2Q9sG) as equation #4.

I'm really confused as to how this gets converted into the following code:

epistemic = np.mean(p_hat**2, axis=0) - np.mean(p_hat, axis=0)**2
aleatoric = np.mean(p_hat*(1-p_hat), axis=0)

I think I may be confused due to the symbols in the equation though, for instance "diag" and the circle with a cross in it. How is the diag and circle represented as such in Python?

Thanks.

Answer 1

The linked article explains below eq. (2) the definition of this "outer square" operation:

v^⊗2 = v v^T.

In the field of machine learning, vectors are to be interpreted as column vectors. In numpy, if v is a vector of shape (n,), you would write:

v.reshape(-1, 1) * v

The result is a shape (n, n) array. If you want to stay closer to the notation with column and row vectors, you could also write:

v_column = v.reshape(-1, 1)
result = v_column @ v_column.T

The function diag is the same as np.diag: you feed it a vector and you get a diagonal matrix.

How the Python snippet in your question implements the equations from the paper is hard to tell without information about what the shape of p_hat is and which axis represents t in the equation. In the most plausible definition, p_hat would have shape (T, n), in which case np.mean(p_hat**2, axis=0) would return shape (n,), which is not consistent with the equation that you quoted from the paper that should result in an (n, n) shape.

Given that the epistemic and aleatoric uncertainties in the paper seem to be scalars rather than vectors, I suspect that the authors made an error in the definition of the ⊗2 exponent: they should have written

v^⊗2 = v^T v

which translates to Python (on shape (n,) arrays) as

np.sum(v**2)