code: def restore_from_path(ckpt_dir: str) -> CheckpointState:
      """Restore the state stored in ckpt_dir."""
      array_path = os.path.join(ckpt_dir, _ARRAY_NAME)
      exemplar_path = os.path.join(ckpt_dir, _EXEMPLAR_NAME)
  
      with open(exemplar_path, "rb") as f:
          exemplar = pickle.load(f)
  
      with open(array_path, "rb") as f:
          files = np.load(f, allow_pickle=True)
          flat_state = [files[key] for key in files.files]
      unflattened_tree = tree.unflatten_as(exemplar, flat_state)
  
      def maybe_convert_to_python(value, numpy):
          return value if numpy else value.item()
  
      return tree.map_structure(
          maybe_convert_to_python, unflattened_tree, exemplar
      )
  
file_path: acme/acme/jax/savers.py
imports: import datetime
  import os
  import pickle
  from typing import Any
  from absl import logging
  from acme import core
  from acme.tf import savers
  import jax
  import numpy
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: restore_from_path(ckpt_dir) and returns: tree, unknown