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: Calls the sharding callback with shardable_tensors.
  
  Args:
    sharding_callback: ShardingCallback. The callback function wrapper that
      splits shardable_tensors into shards.
  
  Returns:
    A list of shards.
docstr_tok: ["Calls","the","sharding","callback","with","shardable_tensors",".","Args",":","sharding_callback",":","ShardingCallback",".","The","callback","function","wrapper","that","splits","shardable_tensors","into","shards",".","Returns",":","A","list","of","shards","."]
code: def _get_shards_by_task(
      self, sharding_callback: sharding_util.ShardingCallback
  ) -> Sequence[sharding_util.TensorSliceDict]:
      """Calls the sharding callback with shardable_tensors.
  
      Args:
        sharding_callback: ShardingCallback. The callback function wrapper that
          splits shardable_tensors into shards.
  
      Returns:
        A list of shards.
      """
      shardable_tensors_by_task = {}
      for shardable_tensor in self._shardable_tensors:
          tensor_val = shardable_tensor.tensor
          tensor_shape = shardable_tensor.shape
          save_spec = (
              shardable_tensor._tensor_save_spec
          )  # pylint: disable=protected-access
          with ops.device(shardable_tensor.device):
              save_spec_tensor = save_spec.tensor
  
          if tensor_val is None and save_spec_tensor is None:
              # A tensor value of `None` indicates that this SaveableObject gets
              # recorded in the object graph, but that no value is saved in the
              # checkpoint.
              continue
          elif save_spec_tensor is not None:
              # Pull the tensor value from _tensor_save_spec.
              tensor_val = save_spec_tensor
              tensor_shape = save_spec_tensor.shape
  
              # Propagate the save spec name and/or slice spec when they are tensors.
              # This makes sure properties like `layout` for dtensor names/slice specs
              # are preserved during sharding.
              if isinstance(save_spec.name, tensor_lib.Tensor):
                  tensor_val._wrapped_name = (
                      save_spec.name
                  )  # pylint: disable=protected-access
              if isinstance(
                  shardable_tensor.slice_spec, tensor_lib.Tensor
              ):
                  tensor_val._wrapped_slice_spec = (
                      save_spec.slice_spec
                  )  # pylint: disable=protected-access
  
          task = device_lib.DeviceSpec.from_string(
              saveable_object_util.set_cpu0(
                  shardable_tensor.device.to_string()
              )
          )
          shardable_tensors_by_task.setdefault(task, []).append(
              dataclasses.replace(
                  shardable_tensor,
                  tensor=tensor_val,
                  shape=tensor_shape,
              )
          )
  
      sharding_callback = (
          sharding_callback or sharding_policies.ShardByTaskPolicy()
      )
      shards_by_task = [
          (task, sharding_callback(shardable_tensors))
          for task, shardable_tensors in shardable_tensors_by_task.items()
      ]
      return shards_by_task
  
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 6 for calling the MultiDeviceSaver obj's underlying member method code with example usage: obj._get_shards_by_task(sharding_callback) and returns: shards_by_task