class_docstr: Policy that splits tensors into shards with a max shard size.
docstr: Callback to split tensors into shards with a max shard size.
  
  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","with","a","max","shard","size",".","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 with a max shard size.
  
      Args:
        shardable_tensors: A list of ShardableTensors.
  
      Returns:
        List of shard dicts containing tensors.
            [ {checkpoint key: {slice_spec: tensor} } ]
      """
      tensors_by_shard = []
      large_scalars = []
  
      shard_size_remaining = self.max_shard_size
      for shardable_tensor in shardable_tensors:
          root_tensor = shardable_tensor.tensor
          root_shape = shardable_tensor.shape
          dtype = shardable_tensor.dtype
          checkpoint_key = shardable_tensor.checkpoint_key
  
          dtype_size = dtypes.as_dtype(dtype).size
          total_size = (
              root_shape.num_elements() * dtype_size
          )  # in bytes
  
          if dtype == dtypes.string:
              if not context.executing_eagerly():
                  with ops.device(shardable_tensor.device):
                      root_tensor = ops.get_default_session().run(
                          root_tensor
                      )
              if root_shape.rank is None or root_shape.rank == 0:
                  sizes = [
                      string_ops.string_length(root_tensor, unit="BYTE")
                  ]
              else:
                  sizes = [
                      string_ops.string_length(elem, unit="BYTE")
                      for elem in root_tensor
                  ]
              if context.executing_eagerly():
                  sizes = [size.numpy() for size in sizes]
              else:
                  with ops.device(shardable_tensor.device):
                      sizes = ops.get_default_session().run(sizes)
              total_size = sum(sizes)
              dtype_size = max(sizes)
  
          if total_size > self.max_shard_size and (
              root_shape.rank is None or root_shape.rank == 0
          ):
              logging.warning(
                  "Tensor %s is a scalar of size %s bytes and cannot be "
                  "partitioned into a shard of max shard size %s bytes. "
                  "It will be added as an individual shard that exceeds "
                  "the max shard size.",
                  checkpoint_key,
                  total_size,
                  self.max_shard_size,
              )
              large_scalars.append(
                  {
                      checkpoint_key: {
                          shardable_tensor.slice_spec: root_tensor
                      }
                  }
              )
              continue
  
          # Partition tensor and add partitions to shards.
          working_tensor = root_tensor
          working_tensor_var_offset = [0] * root_shape.rank
          working_tensor_shape = root_shape
          working_tensor_size = total_size
          while working_tensor_size > shard_size_remaining:
              part_axis_and_size = self._get_next_partition(
                  shard_size_remaining=shard_size_remaining,
                  shape=working_tensor_shape,
                  dtype_size=dtype_size,
                  num_elems=working_tensor_shape.num_elements(),
              )
  
              (
                  remaining_size,
                  (remaining_offset, remaining_shape),
              ) = self._add_partition(
                  root_shardable_tensor=shardable_tensor,
                  dtype_size=dtype_size,
                  working_tensor_offset=working_tensor_var_offset,
                  part_axis_and_size=part_axis_and_size,
                  shard_size_remaining=shard_size_remaining,
                  max_shard_size=self.max_shard_size,
                  tensors_by_shard=tensors_by_shard,
                  large_scalars=large_scalars,
              )
  
              if remaining_size is None:
                  # Tensor partition couldn't fit in remaining shard space. Try again
                  # with the next full shard.
                  tensors_by_shard.append({})
                  shard_size_remaining = self.max_shard_size
              else:
                  working_tensor = array_ops.slice(
                      root_tensor,
                      begin=remaining_offset,
                      size=remaining_shape,
                  )
                  working_tensor_var_offset = remaining_offset
                  working_tensor_shape = working_tensor.shape
                  working_tensor_size = (
                      int(math.prod(remaining_shape)) * dtype_size
                  )
                  shard_size_remaining = remaining_size
  
          if working_tensor_shape.num_elements() > 0:
              remaining_tensor_slice_spec = (
                  variables.Variable.SaveSliceInfo(
                      full_name=checkpoint_key,
                      full_shape=root_shape,
                      var_offset=working_tensor_var_offset,
                      var_shape=working_tensor_shape,
                  ).spec.strip()
              )
              if not tensors_by_shard:
                  tensors_by_shard.append({})
              (
                  tensors_by_shard[-1].setdefault(checkpoint_key, {})[
                      remaining_tensor_slice_spec
                  ]
              ) = working_tensor
          shard_size_remaining -= working_tensor_size
  
      return tensors_by_shard + large_scalars
  
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
label_desc: Use image node_id 5 for calling the MaxShardSizePolicy obj's underlying member method code with example usage: obj.__call__(shardable_tensors) and returns: unknown