Merge pull request #2592 from IvyZX:IvyZX-patch-1

PiperOrigin-RevId: 486754847

Author: Flax Authors

flax/training/checkpoints.py

@@ -231,7 +231,7 @@ def _check_mpa_errors():
       raise errors.MPARestoreTargetRequiredError(ckpt_path, step)
 
   def _safe_deserialize(
-      target_mpas: List[Tuple[Tuple[Any], MultiprocessArrayType, str]],
+      target_mpas: List[Tuple[Tuple[Any, ...], MultiprocessArrayType, str]],
       gda_manager: GlobalAsyncCheckpointManager) -> List[MultiprocessArrayType]:
     gda_manager.wait_until_finished()