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: Save the saveable objects to a checkpoint with `file_prefix`.
  
  Args:
    file_prefix: A string or scalar string Tensor containing the prefix to
      save under.
    options: Optional `CheckpointOptions` object.
  Returns:
    An `Operation`, or None when executing eagerly.
docstr_tok: ["Save","the","saveable","objects","to","a","checkpoint","with","`","file_prefix","`",".","Args",":","file_prefix",":","A","string","or","scalar","string","Tensor","containing","the","prefix","to","save","under",".","options",":","Optional","`","CheckpointOptions","`","object",".","Returns",":","An","`","Operation","`",",","or","None","when","executing","eagerly","."]
code: def save(
      self,
      file_prefix: tensor_lib.Tensor,
      options: "checkpoint_options.CheckpointOptions | None" = None,
  ) -> ops.Operation:
      """Save the saveable objects to a checkpoint with `file_prefix`.
  
      Args:
        file_prefix: A string or scalar string Tensor containing the prefix to
          save under.
        options: Optional `CheckpointOptions` object.
      Returns:
        An `Operation`, or None when executing eagerly.
      """
      options = options or checkpoint_options.CheckpointOptions()
  
      # IMPLEMENTATION DETAILS: most clients should skip.
      #
      # Suffix for any well-formed "checkpoint_prefix", when sharded.
      # Transformations:
      # * Users pass in "save_path" in save() and restore().  Say "myckpt".
      # * checkpoint_prefix gets fed <save_path><sharded_suffix>.
      #
      # Example:
      #   During runtime, a temporary directory is first created, which contains
      #   files
      #
      #     <train dir>/myckpt_temp/
      #        part-?????-of-?????{.index, .data-00000-of-00001}
      #
      #   Before .save() finishes, they will be (hopefully, atomically) renamed to
      #
      #     <train dir>/
      #        myckpt{.index, .data-?????-of-?????}
      #
      #   Filesystems with eventual consistency (such as S3), don't need a
      #   temporary location. Using a temporary directory in those cases might
      #   cause situations where files are not available during copy.
      #
      # Users only need to interact with the user-specified prefix, which is
      # "<train dir>/myckpt" in this case.  Save() and Restore() work with the
      # prefix directly, instead of any physical pathname.  (On failure and
      # subsequent restore, an outdated and orphaned temporary directory can be
      # safely removed.)
      with ops.device("CPU"):
          sharded_suffix = array_ops.where(
              string_ops.regex_full_match(file_prefix, "^s3://.*"),
              constant_op.constant(".part"),
              constant_op.constant("_temp/part"),
          )
          tmp_checkpoint_prefix = string_ops.string_join(
              [file_prefix, sharded_suffix]
          )
          registered_paths = {
              saver_name: registered_saver_filename(
                  file_prefix, saver_name
              )
              for saver_name in self._registered_savers
          }
  
      def save_fn() -> ops.Operation:
          saved_prefixes = []
          # Save with the registered savers. These run before default savers due to
          # the API contract.
          for saver_name, (
              save_fn,
              _,
          ) in self._registered_savers.items():
              maybe_saved_prefixes = save_fn(
                  registered_paths[saver_name]
              )
              if maybe_saved_prefixes is not None:
                  flattened_saved_prefixes = nest.flatten(
                      maybe_saved_prefixes
                  )
                  if not all(
                      tensor_util.is_tf_type(x)
                      and x.dtype == dtypes.string
                      for x in flattened_saved_prefixes
                  ):
                      raise ValueError(
                          "Registered saver must return a (maybe empty) list of "
                          f"string type tensors. Got {maybe_saved_prefixes}."
                      )
                  saved_prefixes.extend(flattened_saved_prefixes)
  
          shards_by_task = self._get_shards_by_task(
              options.experimental_sharding_callback
          )
          num_shards_tensor = constant_op.constant(
              sum([len(shards) for _, shards in shards_by_task]),
              name="num_shards",
          )
          sharded_saves = []
  
          shard_idx = 0
          for task, shards in shards_by_task:
              for shard in shards:
                  with ops.device(task):
                      shard_prefix = sharded_filename(
                          tmp_checkpoint_prefix,
                          shard_idx,
                          num_shards_tensor,
                      )
                      shard_idx += 1
                  saved_prefixes.append(shard_prefix)
                  sharded_saves.append(
                      _single_shard_save(
                          shard_prefix, shard, task, options
                      )
                  )
  
          with ops.control_dependencies(sharded_saves):
              # Merge on the io_device if specified, otherwise co-locates the merge op
              # with the last device used.
              tensor_device_spec = self._shardable_tensors[-1].device
              merge_device_spec = (
                  options.experimental_io_device
                  or saveable_object_util.set_cpu0(
                      tensor_device_spec.to_string()
                  )
              )
              with ops.device(merge_device_spec):
                  # V2 format write path consists of a metadata merge step.  Once
                  # merged, attempts to delete the temporary directory,
                  # "<user-fed prefix>_temp".
                  return gen_io_ops.merge_v2_checkpoints(
                      saved_prefixes, file_prefix, delete_old_dirs=True
                  )
  
      # Since this will causes a function re-trace on each save, limit this to the
      # cases where it is needed: eager and when there are multiple tasks. Note
      # that the retrace is needed to ensure we pickup the latest values of
      # options like experimental_io_device.
      if context.executing_eagerly() and self._num_unique_tasks > 1:
          # Explicitly place the identity op on the first device.
          @def_function.function(jit_compile=False)
          def tf_function_save() -> None:
              save_fn()
  
          tf_function_save()
      else:
          return save_fn()
  
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 7 for calling the MultiDeviceSaver obj's underlying member method code with example usage: obj.save(file_prefix, options) and returns: save_fn, gen_io_ops