Support exporting trieste models using tf.saved_model. Note that we don't inherit from tf.Module for two reasons:
it would limit type checking on models (since mypy treats tf.Module as Any)
it wouldn't be enough by at itself anyway as some of our wrappers include other components that are themselves not tf.Modules
Instead we add a get_module_with_variables method for all models, which returns a tf.Module containing references to all the variables contained within the model (using recursive traversal). This can then be extended with any tf,function-decorated methods we wish to export.
Support exporting trieste models using tf.saved_model. Note that we don't inherit from tf.Module for two reasons:
Instead we add a
get_module_with_variablesmethod for all models, which returns a tf.Module containing references to all the variables contained within the model (using recursive traversal). This can then be extended with any tf,function-decorated methods we wish to export.