Test DeepEnsemble with float32

secondmind-labs / trieste

A Bayesian optimization toolbox built on TensorFlow

Apache License 2.0

219 stars 42 forks source link

Test DeepEnsemble with float32 #826

Closed uri-granta closed 7 months ago

uri-granta commented 7 months ago

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)

uri-granta commented 7 months ago

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).

uri-granta commented 7 months ago

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!