implemented predict_y and predict_noise

[hstojic]

this is an important change, it will affect DE model quite a bit, so its important to test downstream consequences

• predict method now outputs estimation uncertainty, which is more appropriate
• DE now supports predict_y method which outputs combined uncertainty
• there is a new method predict_noise that gives mean and variance of the noise (i.e. aleatoric uncertainty), which comes handy in some situations

note that trajectory sampler behaviour will be affected as well, but only in the diversify mode

[ChrisMorter]

Can this be changed to not drop information about the aleatoric uncertainty? Currently Secondmind's calibration sampling logic uses IndependentReparametrizationSampler, which generates samples based upon predictions made using predict, which in turn calls through to predict_encoded.

For this to be suitable for query point generation in active learning, we need to know both the epistemic and aleatoric uncertainty. Hypothetically, we could not use IndependentReparametrizationSampler, and instead write a new sampler which uses predict_ensemble_encoded rather than predict, because that does return information about the aleatoric uncertainty. However, this would be a significant change in the calibration product.

An implementation could look something like:

        ensemble_means, ensemble_vars = self.predict_ensemble_encoded(query_points)
        predicted_means = tf.math.reduce_mean(ensemble_means, axis=-3)
        epistemic_variance = tf.math.reduce_variance(ensemble_means, axis=-3)
        aleatoric_variance = tf.math.reduce_mean(ensemble_vars, axis=-3)
        aleatoric_variance_var = tf.math.reduce_variance(ensemble_vars, axis=-3)

        means = tf.concat([predicted_means, aleatoric_variance], axis=-1)
        vars = tf.concat([epistemic_variance, aleatoric_variance_var], axis=-1)
        
        return means, vars

[ChrisMorter]

Don't let this issue block this PR from being merged though, I appreciate that's likely a difficult change to make since it requires updating all usages, and would be a breaking change for anyone using DeepEnsemble.predict. We can discuss further how to integrate DeepEnsemble for active learning in a future PR.

[hstojic]

I spoke to @uri-granta about this, and its not clear that this usage fits well the interfaces - we decided to proceed as is for now and think through how to deal with this kind of use case

[uri-granta]

LGTM! Just one comment. Also, as discussed, we should avoid changing predict to include the aleatoric uncertainty for now, though there may be a number of different ways we could do this in the future without breaking existing usages.

[uri-granta]

This should have a docstring too, as it's the external API. And also a check_shapes, probably:

@check_shapes(
    "query_points: [broadcast batch..., D]",
    "return[0]: [batch..., E...]",
    "return[1]: [batch..., E...]",
)

Also, does this method generalise to other models beyond DeepEnsemble? If so, we could preemptively define SupportsPredictNoise and EncodedSupportsPredictNoise protocols, just like [Encoded]SupportsPredictY.

[hstojic]

other models in trieste don't have it, so probably not needed at the moment

[hstojic]

added doc and checkshapes otherwise

[uri-granta]

Suggested change

	``query_points``.
	``query_points``, including noise contributions.

[avullo]

lgtm!