I think this is a lucky demo. When I chenge the data generation code, the optimization will be guided wrongly and the variance prediction is wrong.
So I think this uncertainty method only works at a situation that: the value of the diff is close to precsion and log variance. If their value are not at the same scale, the method will broke.
X = np.random.randn(N, Q)
w1 = 2.*1e2
b1 = 8.*1e2
sigma1 = 10 # ground truth
Y1 = X.dot(w1) + b1 + sigma1 * np.random.randn(N, D1)
w2 = 3*1e2
b2 = 3*1e2
sigma2 = 1*1e2 # ground truth
Y2 = X.dot(w2) + b2 + sigma2 * np.random.randn(N, D2)
return X, Y1, Y2