Test DeepEnsemble with float32

[uri-granta]

Related issue(s)/PRs:

Summary

...

Fully backwards compatible: yes / no

PR checklist

• The quality checks are all passing
• The bug case / new feature is covered by tests
• Any new features are well-documented (in docstrings or notebooks)

[hstojic]

a step in the right direction :)
left some comments - this checks if all runs, which is necessary, but certain parts of the pipeline might nevertheless cast back and forth which we don't test here - I'm not sure whats the best way to do that, just saying this is not sufficient - we'll probably need to test each individual component, or perhaps each individual object should have a dtype property which could facilitate that...

cannot recall now, but do we have dtype checks in unit tests of the DE model - e.g. checking prediction, sample methods, and trajectory sampler?

[uri-granta]

By patching tf.cast I've verified that there are no unexpected casts happening in the test. The only casts are in sample_model_index (which casts deep ensemble size to tf.int32) and write_summary_data_based_metrics (which casts dataset observations to the model prediction dtype).

[hstojic]

the diff became a bit mangled, but I think its good now

[uri-granta]

the diff became a bit mangled, but I think its good now

FYI you can add ?w=1 to the diff url to get it to ignore whitespace, which makes it much more readable!