ServingConfig.extra_trackable_resources is now supported in the Orb…

…ax Model path when `support_tf_resources` is `True`.

PiperOrigin-RevId: 699239199

export/orbax/export/modules/obm_module.py

@@ -181,6 +181,14 @@ def _convert_jax_functions_to_obm_functions(
     """Converts the JAX functions to OrbaxModel functions."""
     if serving_config.input_signature is None:
       raise ValueError('serving_config.input_signature is required.')
+    if (
+        not support_tf_resources
+        and serving_config.extra_trackable_resources is not None
+    ):
+      raise ValueError(
+          'serving_config.extra_trackable_resources can only be set when'
+          ' support_tf_resources is True.'
+      )
 
   def _maybe_set_orbax_checkpoint_path(self, jax2obm_kwargs):
     if constants.CHECKPOINT_PATH not in jax2obm_kwargs:

export/orbax/export/modules/obm_module_test.py

@@ -19,6 +19,7 @@
 from orbax.export import constants
 from orbax.export import serving_config as osc
 from orbax.export.modules import obm_module
+import tensorflow as tf
 
 
 class ObmModuleTest(parameterized.TestCase):