Internal

PiperOrigin-RevId: 970142718

Author: hejiang0116

export/orbax/export/constants.py

@@ -95,6 +95,9 @@ class ExportModelType(enum.Enum):
 # Jax2obm_kwargs key for the loader type.
 LOADER_TYPE = 'loader_type'
 
+# SaveOptions extra_options key for TensorFlow SavedModel save options.
+TF_SAVE_OPTIONS = 'tf_save_options'
+
 ################################################################################
 # Proto field names
 ################################################################################

export/orbax/export/export_manager.py

@@ -118,15 +118,17 @@ def serving_signatures(self) -> Mapping[str, Callable[..., Any]]:
   def save(
       self,
       model_path: str,
-      save_options: Optional[tf.saved_model.SaveOptions] = None,
+      save_options: Optional[
+          tf.saved_model.SaveOptions
+      ]=None,
       signature_overrides: Optional[Mapping[str, Callable[..., Any]]] = None,
   ):
     """Saves the JAX model to a Savemodel.
 
     Args:
       model_path: a directory in which to write the SavedModel.
-      save_options: an optional tf.saved_model.SaveOptions for configuring save
-        options.
+      save_options: an optional tf.saved_model.SaveOptions or obm.SaveOptions for
+        configuring save options.
       signature_overrides: signatures to override the self-maintained ones, or
         additional signatures to export.
     """

export/orbax/export/export_manager_obm_test.py

@@ -15,6 +15,7 @@
 """Tests for the `version==ORBAX_MODEL` path of `ExportManager`."""
 
 import os
+from unittest import mock
 
 from absl.testing import absltest
 from absl.testing import parameterized