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
docstr_tok: ["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
      """
      self._shardable_tensors: MutableSequence[
          sharding_util.ShardableTensor
      ] = []
      # Keep these two data structures so that we can map restored tensors to
      # the Trackable restore functions.
      self._keys_to_restore_fn: MutableMapping[
          sharding_util.TensorSlice,
          Callable[Mapping[str, tensor_lib.Tensor]],
      ] = {}
      self._restore_fn_to_keys: MutableMapping[
          Callable[Mapping[str, tensor_lib.Tensor]],
          MutableSequence[sharding_util.TensorSlice],
      ] = {}
  
      unique_tasks = set()
      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, tensor_slice_dict in tensor_dict.items():
              if not isinstance(tensor_slice_dict, dict):
                  # Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
                  tensor_slice_dict = {"": tensor_slice_dict}
  
              for (
                  slice_spec,
                  tensor_save_spec,
              ) in tensor_slice_dict.items():
                  tensor_value = None
                  if not isinstance(
                      tensor_save_spec, saveable_object.SaveSpec
                  ):
                      tensor_value = tensor_save_spec
                      tensor_save_spec = saveable_object.SaveSpec(
                          tensor=tensor_value,
                          slice_spec=slice_spec,
                          name=checkpoint_key,
                          dtype=tensor_save_spec.dtype,
                          device=tensor_save_spec.device,
                      )
  
                  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))
  
                  device = (
                      device_lib.DeviceSpec.from_string(
                          tensor_save_spec.device
                      )
                      if isinstance(tensor_save_spec.device, str)
                      else tensor_save_spec.device
                  )
                  self._shardable_tensors.append(
                      sharding_util.ShardableTensor(
                          _tensor_save_spec=tensor_save_spec,
                          tensor=tensor_value,
                          dtype=tensor_save_spec.dtype,
                          device=device,
                          name=tensor_save_spec.name,
                          shape=None,
                          slice_spec=slice_spec.strip(),
                          checkpoint_key=checkpoint_key,
                          trackable=obj,
                      )
                  )
                  unique_tasks.add(
                      saveable_object_util.set_cpu0(device.to_string())
                  )
  
      self._num_unique_tasks = len(unique_tasks)
  
      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: 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 1 to create a new MultiDeviceSaver object  with example: obj = MultiDeviceSaver(serialized_tensors, registered_savers, call_with_mapped_captures)