Evaluating uncertainty of model predictions

[vl-dud]

Dear @lululxvi
Let's say we have a trained model that uses dropout (dropout_rate is set in the model's network). Is there a method in deepxde to evaluate uncertainty of the outputs when calling model.predict? It should be something similar to DropoutUncertainty callback. I would just make a custom callback as model.predict(x, callbacks=[evaluate_uncertainty]) but Callback object doesn't know anything about x, only about the model.
I can write something like that

y_pred = []
for _ in range(1000):
    y_pred.append(model._outputs(True, np.asarray(X)))
y_std = np.std(y_pred, axis=0)

but this doesn't seem like a good solution to me.

[lululxvi]

You can implement a new callback similar to DropoutUncertainty