De-duplicate get_ts_context usages and move to ts_utils.

PiperOrigin-RevId: 686540442
google-research · Oct 18, 2024 · d35a3a8 · d35a3a8
1 parent 0ff8254
commit d35a3a8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/t5x/checkpoint_utils.py b/t5x/checkpoint_utils.py
@@ -280,7 +280,9 @@ 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()
+  ts_context = ocp.serialization.ts_utils.get_ts_context(
+      use_ocdbt=is_ocdbt_checkpoint
+  )
 
   def _get_param_info(
       name: str,