class_docstr: A single device saver that places tensors on DTensor Device.
docstr: Saves the saveable objects to a checkpoint with `file_prefix`.
  
  Also query the generated shards from the distributed DTensor SaveV2 ops and
  do a MergeV2 on those. Each op here is backed by a global_barrier to avoid
  racing from multiple clients.
  
  Args:
    file_prefix: A string or scalar string Tensor containing the prefix to
      save under.
    options: Optional `CheckpointOptions` object. This is unused in DTensor.
  
  Returns:
    An `Operation`, or None when executing eagerly.
docstr_tok: ["Saves","the","saveable","objects","to","a","checkpoint","with","`","file_prefix","`",".","Also","query","the","generated","shards","from","the","distributed","DTensor","SaveV2","ops","and","do","a","MergeV2","on","those",".","Each","op","here","is","backed","by","a","global_barrier","to","avoid","racing","from","multiple","clients",".","Args",":","file_prefix",":","A","string","or","scalar","string","Tensor","containing","the","prefix","to","save","under",".","options",":","Optional","`","CheckpointOptions","`","object",".","This","is","unused","in","DTensor",".","Returns",":","An","`","Operation","`",",","or","None","when","executing","eagerly","."]
code: def save(
      self,
      file_prefix: str,
      options: Optional[checkpoint_options.CheckpointOptions] = None,
  ) -> Optional[ops.Operation]:
      """Saves the saveable objects to a checkpoint with `file_prefix`.
  
      Also query the generated shards from the distributed DTensor SaveV2 ops and
      do a MergeV2 on those. Each op here is backed by a global_barrier to avoid
      racing from multiple clients.
  
      Args:
        file_prefix: A string or scalar string Tensor containing the prefix to
          save under.
        options: Optional `CheckpointOptions` object. This is unused in DTensor.
  
      Returns:
        An `Operation`, or None when executing eagerly.
      """
      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
      tensor_names = []
      tensors = []
      tensor_slices = []
      for saveable in self._saveable_objects:
          for spec in saveable.specs:
              tensor = spec.tensor
              # A tensor value of `None` indicates that this SaveableObject gets
              # recorded in the object graph, but that no value is saved in the
              # checkpoint.
              if tensor is not None:
                  if api.device_name() != spec.device:
                      # Some small tensors are placed on CPU0 from save manager and
                      # broadcasted to DTensor mesh, e,g., SaveCounter.
                      tensor = api.pack(
                          [tensor]
                          * self._mesh.host_mesh().num_local_devices(),
                          layout.Layout.replicated(
                              self._mesh.host_mesh(),
                              rank=tensor.shape.rank,
                          ),
                      )
                  tensor_names.append(spec.name)
                  tensors.append(tensor)
                  tensor_slices.append(spec.slice_spec)
      return save_restore.sharded_save(
          self._mesh, file_prefix, tensor_names, tensor_slices, tensors
      )
  
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 2 for calling the _DSaver obj's underlying member method code with example usage: obj.save(file_prefix, options) and returns: save_restore