ExportManager#
Manage the exporting of a JAXModule.
ExportManager#
- class orbax.export.export_manager.ExportManager(module, serving_configs)[source][source]#
Exports a JAXModule with pre- and post-processors.
- __init__(module, serving_configs)[source][source]#
ExportManager constructor.
- Parameters:
module (
JaxModule) – the JaxModule to be exported.serving_configs (
Sequence[ServingConfig]) – a sequence of which each element is a ServingConfig cooresponding to a serving signature of the exported SavedModel.
- property tf_module: Module#
Returns the tf.module maintained by the export manager.
- Return type:
Module
- property serving_signatures: Mapping[str, Callable[[...], Any]]#
Returns a map of signature keys to serving functions.
- Return type:
Mapping[str,Callable[…,Any]]
- save(model_path, save_options=None, signature_overrides=None)[source][source]#
Saves the JAX model to a Savemodel.
- Parameters:
model_path (
str) – a directory in which to write the SavedModel.save_options (
Optional[SaveOptions]) – an optional tf.saved_model.SaveOptions for configuring save options.signature_overrides (
Optional[Mapping[str,Callable[…,Any]]]) – signatures to override the self-maintained ones, or additional signatures to export.