Make jax_module optional in ObmExport.

PiperOrigin-RevId: 825586752

Author: hejiang0116

export/orbax/export/export_manager.py

@@ -35,7 +35,7 @@ class ExportManager:
 
   def __init__(
       self,
-      module: jax_module.JaxModule,
+      module: jax_module.JaxModule | None,
       serving_configs: Sequence[osc.ServingConfig],
   ):
     """ExportManager constructor.
@@ -45,9 +45,12 @@ def __init__(
       serving_configs: a sequence of which each element is a `ServingConfig`
         cooresponding to a serving signature of the exported SavedModel.
     """
-    self._version = module.export_version
     self._jax_module = module
-    if self._version == constants.ExportModelType.ORBAX_MODEL:
+    if (
+        not self._jax_module
+        or self._jax_module.export_version
+        == constants.ExportModelType.ORBAX_MODEL
+    ):
       self._serialization_functions = obm_export.ObmExport(
           self._jax_module, serving_configs
       )
@@ -59,7 +62,11 @@ def __init__(
   @property
   def tf_module(self) -> tf.Module:
     """Returns the tf.module maintained by the export manager."""
-    if self._version == constants.ExportModelType.ORBAX_MODEL:
+    if (
+        not self._jax_module
+        or self._jax_module.export_version
+        == constants.ExportModelType.ORBAX_MODEL
+    ):
       raise TypeError(
           'tf_module is not implemented for export version'
           ' ExportModelType.ORBAX_MODEL.'

export/orbax/export/export_manager_obm_test.py

@@ -18,10 +18,12 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
+
 from orbax.export import constants
 from orbax.export import export_manager
 from orbax.export import export_testing_utils
 from orbax.export import oex_orchestration
+from orbax.export import serving_config as sc
 import tensorflow as tf
 
 

export/orbax/export/obm_export.py

@@ -43,7 +43,7 @@ class ObmExport(export_base.ExportBase):
 
   def __init__(
       self,
-      module: jax_module.JaxModule,
+      module: jax_module.JaxModule | None,
       serving_configs: Sequence[osc.ServingConfig],
   ):
     """Initializes the ObmExport class."""