class_docstr: Saves checkpoints directly from multiple devices.
  
  Note that this is a low-level utility which stores Tensors in the keys
  specified by `SaveableObject`s. Higher-level utilities for object-based
  checkpointing are built on top of it.
docstr: 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.
    options: Optional `CheckpointOptions` object.
  
  Returns:
    When not run eagerly or when saving on a single device, returns a
    dictionary mapping from SaveableObject names to restore operations;
    otherwise, returns an empty dict.
code: def restore(
      self,
      file_prefix: tensor_lib.Tensor,
      options: "checkpoint_options.CheckpointOptions | None" = None,
  ) -> Mapping[str, ops.Operation]:
      """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.
        options: Optional `CheckpointOptions` object.
  
      Returns:
        When not run eagerly or when saving on a single device, returns a
        dictionary mapping from SaveableObject names to restore operations;
        otherwise, returns an empty dict.
      """
      options = options or checkpoint_options.CheckpointOptions()
  
      def restore_fn() -> Mapping[str, ops.Operation]:
          restore_fn_inputs = {}
          restore_fn_input_count = {
              fn: len(keys)
              for fn, keys in self._restore_fn_to_keys.items()
          }
  
          restore_ops = {}
          # Sort by device_spec to avoid propagating non-deterministic dictionary
          # ordering in some Python versions.
          for device_spec, tensor_slice_dict in sorted(
              self._tensors_by_task.items()
          ):
              with ops.device(device_spec):
                  # Load values from checkpoint
                  restored_tensor_dict = _single_task_restore(
                      file_prefix, tensor_slice_dict, options
                  )
  
                  # Map restored tensors to the corresponding restore_fn, and see if all
                  # inputs have all been loaded. Call `restore_fn` if that is the case.
                  for (
                      checkpoint_key,
                      slice_and_tensor,
                  ) in restored_tensor_dict.items():
                      for (
                          slice_spec,
                          tensor,
                      ) in slice_and_tensor.items():
                          restore_fn = self._keys_to_restore_fn[
                              (checkpoint_key, slice_spec)
                          ]
  
                          # Processing the returned restored_tensor_dict to prepare for the
                          # Trackable `restore` function. The `restore` function expects a
                          # map of `string name (checkpoint_key) -> Tensor`. Unless there is
                          # a slice_spec, in which case the map will be of
                          # `string name (checkpoint_key)-> slice_spec -> Tensor`.
                          if slice_spec:
                              (
                                  restore_fn_inputs.setdefault(
                                      restore_fn, {}
                                  ).setdefault(checkpoint_key, {})[
                                      slice_spec
                                  ]
                              ) = tensor
                          else:
                              restore_fn_inputs.setdefault(
                                  restore_fn, {}
                              )[checkpoint_key] = tensor
                          restore_fn_input_count[restore_fn] -= 1
  
                          if restore_fn_input_count[restore_fn] == 0:
                              restored_tensors = {}
                              # Extracts the substring after the "/.ATTRIBUTES/" in the
                              # ckpt_key from restore_fn_inputs[restore_fn] to
                              # restored_tensors. For example, if restore_fn_input[restore_fn]
                              # is dict { "/.ATTIBUTES/a": Tensor}, restored_tensors will be
                              # changed to dict {"a": Tensor}
                              for ckpt_key, tensor in restore_fn_inputs[
                                  restore_fn
                              ].items():
                                  restored_tensors[
                                      trackable_utils.extract_local_name(
                                          ckpt_key
                                      )
                                  ] = tensor
                              ret = restore_fn(restored_tensors)
                              if isinstance(ret, dict):
                                  restore_ops.update(ret)
          # Run registered restore methods after the default restore ops.
          for _, (_, restore_fn) in self._registered_savers.items():
              restore_fn(file_prefix)
          return restore_ops
  
      has_custom_device_saver = any(
          [context.is_custom_device(ds) for ds in self._tensors_by_task]
      )
      # Since this will cause a function re-trace on each restore, limit this to
      # cases where it is needed: eager and when there are multiple tasks or any
      # device_spec is a custom device. Note that the retrace is needed to ensure
      # we pickup the latest values of options like experimental_io_device.
      #
      # We run in a function when there is a custom device saver because custom
      # devices, such as DTensor, usually do a sharded save and restore.
      # Doing a sharded save and restore requires knowledge about what shards
      # of variables we are restoring to. In practice, this means that custom
      # devices need the AssignVariableOps along with the Restore op within the
      # same graph to infer shapes and shard specs for Restore op.
      if context.executing_eagerly() and (
          len(self._tensors_by_task) > 1 or has_custom_device_saver
      ):
  
          @def_function.function(jit_compile=False, autograph=False)
          def tf_function_restore() -> Mapping[str, ops.Operation]:
              restore_fn()
              return {}
  
          restore_ops = tf_function_restore()
      else:
          restore_ops = restore_fn()
  
      return restore_ops
  
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 7 for calling the MultiDeviceSaver obj's underlying member method code with example usage: obj.restore(file_prefix, options) and returns: restore_ops, restore_ops, dict