Fix local restore by re-mapping device ids directly instead of inferr…

…ing them from how process indexes changed across restarts with some false assumptions.

PiperOrigin-RevId: 700737164

MaxText/max_utils.py

@@ -325,6 +325,7 @@ def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys):
     jax.distributed.initialize()
 
   ocp.multihost.initialize_runtime_to_distributed_ids()
+  ocp.multihost.initialize_distributed_to_device_ids()
 
 
 def _retrieve_jax_init_info(raw_keys):