class_docstr: Policy that splits tensors into shards based on their device spec task.
docstr: Callback to split tensors into shards based on their device spec task.
  
  Args:
    shardable_tensors: A list of ShardableTensors.
  
  Returns:
    List of shard dicts containing tensors.
        [ {checkpoint key: {slice_spec: tensor} } ]
docstr_tok: ["Callback","to","split","tensors","into","shards","based","on","their","device","spec","task",".","Args",":","shardable_tensors",":","A","list","of","ShardableTensors",".","Returns",":","List","of","shard","dicts","containing","tensors",".","[","{","checkpoint","key",":","{","slice_spec",":","tensor","}","}","]"]
code: def __call__(
      self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
  ) -> Sequence[sharding_util.TensorSliceDict]:
      """Callback to split tensors into shards based on their device spec task.
  
      Args:
        shardable_tensors: A list of ShardableTensors.
  
      Returns:
        List of shard dicts containing tensors.
            [ {checkpoint key: {slice_spec: tensor} } ]
      """
      tensors_by_task = {}
  
      for shardable_tensor in shardable_tensors:
          tensor = shardable_tensor.tensor
          checkpoint_key = shardable_tensor.checkpoint_key
          slice_spec = shardable_tensor.slice_spec
  
          (
              tensors_by_task.setdefault(checkpoint_key, {})[slice_spec]
          ) = tensor
  
      return [tensors_by_task]
  
imports: import math
  from typing import MutableSequence, Sequence
  from absl import logging
  from tensorflow.python.checkpoint.sharding import sharding_util
  from tensorflow.python.eager import context
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import string_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.trackable import base
  from tensorflow.python.util import tf_export
label_desc: Use image node_id 2 for calling the ShardByTaskPolicy obj's underlying member method code with example usage: obj.__call__(shardable_tensors) and returns: list