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: Specify a list of `SaveableObject`s to save and restore.
  
  Args:
    serialized_tensors: A dictionary mapping `Trackable` to a tensor dict,
      which maps checkpoint_key -> (slice_spec ->) -> Tensor/SaveSpec. The
      `Trackable` key is used to get the `restore_from_tensors` function,
      and may be `None` if the tensor is not meant to be restored.
    registered_savers: A dictionary mapping `registration.RegisteredSaver`
      namedtuples to a dictionary of named Trackables. The keys of the
      Trackable dictionary are string names that uniquely identify the
      Trackable in the checkpoint.
    call_with_mapped_captures: TODO
code: def __init__(
      self,
      serialized_tensors: Mapping[
          base.Trackable, sharding_util.TensorSliceDict
      ],
      registered_savers: "RegisteredSaversDict | None" = None,
      call_with_mapped_captures: "MappedCapturesCallable | None" = None,
  ):
      """Specify a list of `SaveableObject`s to save and restore.
  
      Args:
        serialized_tensors: A dictionary mapping `Trackable` to a tensor dict,
          which maps checkpoint_key -> (slice_spec ->) -> Tensor/SaveSpec. The
          `Trackable` key is used to get the `restore_from_tensors` function,
          and may be `None` if the tensor is not meant to be restored.
        registered_savers: A dictionary mapping `registration.RegisteredSaver`
          namedtuples to a dictionary of named Trackables. The keys of the
          Trackable dictionary are string names that uniquely identify the
          Trackable in the checkpoint.
        call_with_mapped_captures: TODO
      """
      # Keep these two data structures so that we can map restored tensors to
      # the Trackable restore functions.
      self._keys_to_restore_fn = {}
      self._restore_fn_to_keys = {}
  
      # Extract serialized tensors and separate by task, which contains 1 or more
      # devices. A DeviceSpec specifies both the task and device.
  
      # device_spec -> checkpoint key -> (slice_spec ->) tensor
      self._tensors_by_task = {}
  
      for obj, tensor_dict in serialized_tensors.items():
          restore_fn = (
              _restore_noop
              if obj is None
              else obj._restore_from_tensors
          )
  
          # Divide tensor_dict by task.
          for checkpoint_key, maybe_tensor in tensor_dict.items():
              if not isinstance(maybe_tensor, dict):
                  # Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
                  maybe_tensor = {"": maybe_tensor}
  
              for slice_spec, tensor in maybe_tensor.items():
                  if (
                      checkpoint_key,
                      slice_spec,
                  ) in self._keys_to_restore_fn:
                      raise ValueError(
                          "Recieved multiple tensors with the same checkpoint key and "
                          "slice spec. This is invalid because one will overwrite the "
                          "other in the checkpoint. This indicates a bug in the "
                          "Checkpoint key-generation."
                      )
                  self._keys_to_restore_fn[
                      (checkpoint_key, slice_spec)
                  ] = restore_fn
                  self._restore_fn_to_keys.setdefault(
                      restore_fn, []
                  ).append((checkpoint_key, slice_spec))
  
                  # By setting the device within the DeviceSpec to CPU:0, we ensure that
                  # we group the tensors by task.
                  tensor_task = saveable_object_util.set_cpu0(
                      tensor.device
                  )
                  (
                      self._tensors_by_task.setdefault(
                          tensor_task, {}
                      ).setdefault(checkpoint_key, {})[slice_spec]
                  ) = tensor
  
      self._registered_savers = {}
      if registered_savers:
          for registered_name, trackables in registered_savers.items():
              save_fn = _get_mapped_registered_save_fn(
                  registration.get_save_function(registered_name),
                  trackables,
                  call_with_mapped_captures,
              )
              restore_fn = _get_mapped_registered_restore_fn(
                  registration.get_restore_function(registered_name),
                  trackables,
                  call_with_mapped_captures,
              )
              self._registered_savers[registered_name] = (
                  save_fn,
                  restore_fn,
              )
  
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 1 to create a new MultiDeviceSaver object  with example: obj = MultiDeviceSaver(serialized_tensors, registered_savers, call_with_mapped_captures)