code: def _single_shard_restore(
      file_prefix: tensor_lib.Tensor,
      shardable_tensors: Sequence[sharding_util.ShardableTensor],
      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.
        shardable_tensors: A list of ShardableTensors to restore.
        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 shardable_tensor in shardable_tensors:
          if (
              shardable_tensor._tensor_save_spec
          ):  # pylint: disable=protected-access
              name = (
                  shardable_tensor._tensor_save_spec.name
              )  # pylint: disable=protected-access
              spec = (
                  shardable_tensor._tensor_save_spec.slice_spec
              )  # pylint: disable=protected-access
          else:
              name, spec = (
                  shardable_tensor.checkpoint_key,
                  shardable_tensor.slice_spec,
              )
          tensor_names.append(name)
          slice_specs.append(spec)
          tensor_dtypes.append(shardable_tensor.dtype)
  
      restore_device = options.experimental_io_device or "cpu:0"
      with ops.device(restore_device):
          restored_tensors = io_ops.restore_v2(
              file_prefix, tensor_names, slice_specs, tensor_dtypes
          )
  
      restored_tensor_dict = {}
      for shardable_tensor in shardable_tensors:
          restored_tensor = restored_tensors.pop(0)
          (
              restored_tensor_dict.setdefault(
                  shardable_tensor.checkpoint_key, {}
              )[shardable_tensor.slice_spec]
          ) = restored_tensor
      return restored_tensor_dict
  
imports: import dataclasses
  from typing import Callable, Mapping, MutableMapping, MutableSequence, Sequence
  from tensorflow.core.protobuf import saver_pb2
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.checkpoint.sharding import sharding_policies
  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 device
  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_shard_restore(file_prefix, shardable_tensors, options) and returns: restored_tensor_dict