class_docstr: Policy that splits tensors into shards with a max shard size.
docstr: Adds the tensor partition to the shard, if possible.
  
  Args:
    root_shardable_tensor: The full tensor being partitioned.
    dtype_size: Size in bytes of the dtype of the working tensor.
    working_tensor_offset: The offset of the working tensor in the full
        tensor.
    part_axis_and_size: A tuple containing the axis of the partition and that
        partition size.
    shard_size_remaining: Size in bytes of the space remaining in the shard.
    max_shard_size: Max size in bytes allowed for a checkpoint shard.
    tensors_by_shard: List of shard dicts containing tensors.
        [ {checkpoint key: {slice_spec: tensor} } ]
    large_scalars: List of shard dicts containing scalars too large to fit in
        the max_shard_size. [ {checkpoint key: {slice_spec: tensor} } ]
  
  Returns:
    A tuple containing the size of the slice that was added to the shard and
        the offset & shape of the remaining portion of the tensor.
docstr_tok: ["Adds","the","tensor","partition","to","the","shard",",","if","possible",".","Args",":","root_shardable_tensor",":","The","full","tensor","being","partitioned",".","dtype_size",":","Size","in","bytes","of","the","dtype","of","the","working","tensor",".","working_tensor_offset",":","The","offset","of","the","working","tensor","in","the","full","tensor",".","part_axis_and_size",":","A","tuple","containing","the","axis","of","the","partition","and","that","partition","size",".","shard_size_remaining",":","Size","in","bytes","of","the","space","remaining","in","the","shard",".","max_shard_size",":","Max","size","in","bytes","allowed","for","a","checkpoint","shard",".","tensors_by_shard",":","List","of","shard","dicts","containing","tensors",".","[","{","checkpoint","key",":","{","slice_spec",":","tensor","}","}","]","large_scalars",":","List","of","shard","dicts","containing","scalars","too","large","to","fit","in","the","max_shard_size",".","[","{","checkpoint","key",":","{","slice_spec",":","tensor","}","}","]","Returns",":","A","tuple","containing","the","size","of","the","slice","that","was","added","to","the","shard","and","the","offset","&","shape","of","the","remaining","portion","of","the","tensor","."]
code: def _add_partition(
      self,
      root_shardable_tensor: sharding_util.ShardableTensor,
      dtype_size: int,
      working_tensor_offset: Sequence[int],
      part_axis_and_size: _PartitionAxisAndSize,
      shard_size_remaining: int,
      max_shard_size: int,
      tensors_by_shard: MutableSequence[sharding_util.TensorSliceDict],
      large_scalars: MutableSequence[sharding_util.TensorSliceDict],
  ) -> tuple[tensor_lib.Tensor, _OffsetAndShape]:
      """Adds the tensor partition to the shard, if possible.
  
      Args:
        root_shardable_tensor: The full tensor being partitioned.
        dtype_size: Size in bytes of the dtype of the working tensor.
        working_tensor_offset: The offset of the working tensor in the full
            tensor.
        part_axis_and_size: A tuple containing the axis of the partition and that
            partition size.
        shard_size_remaining: Size in bytes of the space remaining in the shard.
        max_shard_size: Max size in bytes allowed for a checkpoint shard.
        tensors_by_shard: List of shard dicts containing tensors.
            [ {checkpoint key: {slice_spec: tensor} } ]
        large_scalars: List of shard dicts containing scalars too large to fit in
            the max_shard_size. [ {checkpoint key: {slice_spec: tensor} } ]
  
      Returns:
        A tuple containing the size of the slice that was added to the shard and
            the offset & shape of the remaining portion of the tensor.
      """
      root_tensor = root_shardable_tensor.tensor
      root_tensor_shape = root_shardable_tensor.shape
      checkpoint_key = root_shardable_tensor.checkpoint_key
  
      if root_tensor_shape.rank is None or root_tensor_shape.rank == 0:
          return None, (None, None)
  
      min_axis, part_size = part_axis_and_size
  
      # Add what we can to the current shard.
      slice_offset = working_tensor_offset
      slice_shape = [
          root_tensor_shape[i] - slice_offset[i]
          for i in range(root_tensor_shape.rank)
      ]
      slice_shape[min_axis] = part_size
      slice_size_in_bytes = int(math.prod(slice_shape)) * dtype_size
      with ops.device(root_shardable_tensor.device):
          tensor_slice = array_ops.slice(
              root_tensor, begin=slice_offset, size=slice_shape
          )
      slice_spec = variables.Variable.SaveSliceInfo(
          full_name=checkpoint_key,
          full_shape=root_tensor_shape,
          var_offset=slice_offset,
          var_shape=slice_shape,
      ).spec.strip()
      remaining_size = shard_size_remaining
      if slice_size_in_bytes > max_shard_size:
          logging.warning(
              "Slice %s of 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.",
              slice_spec,
              checkpoint_key,
              slice_size_in_bytes,
              max_shard_size,
          )
          large_scalars.append(
              {checkpoint_key: {slice_spec: tensor_slice}}
          )
      elif slice_size_in_bytes > shard_size_remaining:
          # Smallest partition can't fit in the remaining shard space. Start fresh
          # with a new shard.
          return None, (None, None)
      else:
          if not tensors_by_shard or shard_size_remaining < 1:
              tensors_by_shard.append({})
              remaining_size = max_shard_size
          (
              tensors_by_shard[-1].setdefault(checkpoint_key, {})[
                  slice_spec
              ]
          ) = tensor_slice
          remaining_size -= slice_size_in_bytes
  
      # Get remaining portion of tensor to add to the next shard(s).
      slice_offset[min_axis] += part_size
      slice_shape = [
          root_tensor_shape[i] - slice_offset[i]
          for i in range(root_tensor_shape.rank)
      ]
  
      return (remaining_size, (slice_offset, slice_shape))
  
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 4 for calling the MaxShardSizePolicy obj's underlying member method code with example usage: obj._add_partition(root_shardable_tensor, dtype_size, working_tensor_offset, part_axis_and_size, shard_size_remaining, max_shard_size, tensors_by_shard, large_scalars) and returns: remaining_size,  slice_offset,  slice_shape, None,  None,  None, None,  None,  None