class_docstr: A single device saver that places tensors on DTensor Device.
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. This is unused in DTensor.
  
  Returns:
    A dictionary mapping from SaveableObject names to restore operations.
docstr_tok: ["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",".","This","is","unused","in","DTensor",".","Returns",":","A","dictionary","mapping","from","SaveableObject","names","to","restore","operations","."]
code: def restore(
      self,
      file_prefix: str,
      options: Optional[checkpoint_options.CheckpointOptions] = None,
  ) -> Dict[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. This is unused in DTensor.
  
      Returns:
        A dictionary mapping from SaveableObject names to restore operations.
      """
      if (
          options is not None
          and options.experimental_io_device is not None
      ):
          raise ValueError(
              "Specified experimental_io_device in DTensor checkpoint is not "
              "supported."
          )
      del options
      restore_specs = []
      tensor_structure = []
      for saveable in self._saveable_objects:
          saveable_tensor_structure = []
          tensor_structure.append(saveable_tensor_structure)
          # DTensor change 1 : Gather shapes and layout from original saveable
          # specs.
          # Note that this relies on the fact that the variables are already
          # initialized -- which isn't the behavior we want eventually.
          # TODO(b/159035705): Handle the variable initialization in restore.
          for spec in saveable.specs:
              saveable_tensor_structure.append(spec.name)
              if isinstance(spec, d_variable.DSaveSpec):
                  restore_specs.append(
                      (
                          spec.name,
                          spec.slice_spec,
                          spec.dtype,
                          spec.layout,
                          spec.global_shape,
                      )
                  )
              # Fall back to replicated layouts for non-DTensor saves that constructs
              # normal SaveSpec.
              elif isinstance(spec, saveable_object.SaveSpec):
                  restore_specs.append(
                      (
                          spec.name,
                          spec.slice_spec,
                          spec.dtype,
                          layout.Layout.replicated(
                              self._mesh.host_mesh(),
                              spec.tensor.shape.rank,
                          ).to_string(),
                          spec.tensor.shape.as_list(),
                      )
                  )
      (
          tensor_names,
          tensor_slices,
          tensor_dtypes,
          layouts,
          global_shapes,
      ) = zip(*restore_specs)
      with ops.device(api.device_name()):
          # DTensor change 2 : Run on customized DTensor RestoreV2 op rather than
          # stock TF io_ops.RestoreV2.
          restored_tensors = gen_dtensor_ops.d_tensor_restore_v2(
              prefix=file_prefix,
              tensor_names=tensor_names,
              shape_and_slices=tensor_slices,
              input_shapes=global_shapes,
              input_layouts=layouts,
              dtypes=tensor_dtypes,
          )
      structured_restored_tensors = nest.pack_sequence_as(
          tensor_structure, restored_tensors
      )
      restore_ops = {}
      for saveable, restored_tensors in zip(
          self._saveable_objects, structured_restored_tensors
      ):
          restore_ops[saveable.name] = saveable.restore(
              restored_tensors, restored_shapes=None
          )
      return restore_ops
  
imports: from typing import Dict, List, Optional
  import weakref
  from tensorflow.core.protobuf import trackable_object_graph_pb2
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import gen_dtensor_ops
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import save_restore
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.checkpoint import graph_view
  from tensorflow.python.checkpoint import restore
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.trackable import base
  from tensorflow.python.trackable import data_structures
  from tensorflow.python.training import py_checkpoint_reader
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.training.saving import saveable_object_util
  from tensorflow.python.util import deprecation
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 3 for calling the _DSaver obj's underlying member method code with example usage: obj.restore(file_prefix, options) and returns: restore_ops