code: def _single_task_restore(
      file_prefix: tensor_lib.Tensor,
      tensor_slice_dict: sharding_util.TensorSliceDict,
      options: "checkpoint_options.CheckpointOptions | None" = None,
  ) -> sharding_util.TensorSliceDict:
      """Restore the saveable objects from a checkpoint with `file_prefix`.
  
      Args:
        file_prefix: A string or scalar string Tensor containing the prefix for
          files to read from.
        tensor_slice_dict: A dict mapping checkpoint key -> slice_spec -> tensor.
        options: Optional `CheckpointOptions` object.
  
      Returns:
        A restored tensor dict (maps checkpoint_key -> slice_spec -> tensor).
      """
      options = options or checkpoint_options.CheckpointOptions()
      tensor_names = []
      tensor_dtypes = []
      slice_specs = []
  
      for checkpoint_key, tensor_slices in tensor_slice_dict.items():
          for slice_spec, tensor in tensor_slices.items():
              tensor_dtypes.append(tensor.dtype)
              if isinstance(tensor, saveable_object.SaveSpec):
                  slice_specs.append(tensor.slice_spec)
                  tensor_names.append(tensor.name)
              else:
                  slice_specs.append(slice_spec)
                  tensor_names.append(checkpoint_key)
  
      restore_device_spec = options.experimental_io_device or "cpu:0"
      with ops.device(restore_device_spec):
          restored_tensors = io_ops.restore_v2(
              file_prefix, tensor_names, slice_specs, tensor_dtypes
          )
  
      restored_tensor_dict = {}
      for checkpoint_key, tensor_slices in tensor_slice_dict.items():
          for slice_spec in tensor_slices:
              restored_tensor = restored_tensors.pop(0)
              restored_tensor_dict.setdefault(checkpoint_key, {})[
                  slice_spec
              ] = restored_tensor
      return restored_tensor_dict
  
imports: from typing import Callable, Mapping, Sequence
  from tensorflow.core.protobuf import saver_pb2
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.checkpoint.sharding import sharding_util
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_io_ops
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import string_ops
  from tensorflow.python.saved_model import registration
  from tensorflow.python.trackable import base
  from tensorflow.python.trackable import trackable_utils
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.training.saving import saveable_object_util
  from tensorflow.python.types import core
  from tensorflow.python.util import nest
  from tensorflow.python.util import object_identity
label_desc: Use image node_id 2 for calling a global function with example usage: _single_task_restore(file_prefix, tensor_slice_dict, options) and returns: restored_tensor_dict