De-duplicate get_ts_context usages and move to ts_utils.

PiperOrigin-RevId: 686540442
google-research · Oct 21, 2024 · 6390fa1 · 6390fa1
1 parent b642f30
commit 6390fa1
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/t5x/checkpoint_utils.py b/t5x/checkpoint_utils.py
@@ -280,8 +280,12 @@ def get_restore_parameters(
   restore_args = jax.tree.map(lambda x: ocp.RestoreArgs(), structure)
   flat_param_infos = {}
   is_ocdbt_checkpoint = ocp.type_handlers.is_ocdbt_checkpoint(directory)
-  ts_context = ocp.type_handlers.get_ts_context()
-
+  if hasattr(ocp.serialization, 'ts_utils'):
+    ts_context = ocp.serialization.ts_utils.get_ts_context(
+        use_ocdbt=is_ocdbt_checkpoint
+    )
+  else:
+    ts_context = ocp.type_handlers.get_ts_context()
   def _get_param_info(
       name: str,
       meta_or_value: Union[Any, ocp.metadata.tree.ValueMetadataEntry],