class_docstr: Policy that splits tensors into shards with a max shard size.
  
  Shards may exceed the max shard size if they contain 1. a single scalar/string
  tensor that could not be sliced and exceeds the max shard size or 2. the
  checkpoint object graph, whose size cannot be calculated when saving.
docstr: Gets tensor partition with size closest to shard_size_remaining.
  
  Args:
    shard_size_remaining: Size in bytes of the space remaining in the shard.
    shape: Shape of the working tensor to partition in the remaining
        shard space.
    dtype_size: Size in bytes of the dtype of the working tensor.
    num_elems: Number of elements in the working tensor.
  
  Returns:
    A tuple containing the axis of the next partition and that partition size.
docstr_tok: ["Gets","tensor","partition","with","size","closest","to","shard_size_remaining",".","Args",":","shard_size_remaining",":","Size","in","bytes","of","the","space","remaining","in","the","shard",".","shape",":","Shape","of","the","working","tensor","to","partition","in","the","remaining","shard","space",".","dtype_size",":","Size","in","bytes","of","the","dtype","of","the","working","tensor",".","num_elems",":","Number","of","elements","in","the","working","tensor",".","Returns",":","A","tuple","containing","the","axis","of","the","next","partition","and","that","partition","size","."]
code: def _get_next_partition(
      self,
      shard_size_remaining: int,
      shape: tensor_shape.TensorShape,
      dtype_size: int,
      num_elems: int,
  ) -> _PartitionAxisAndSize:
      """Gets tensor partition with size closest to shard_size_remaining.
  
      Args:
        shard_size_remaining: Size in bytes of the space remaining in the shard.
        shape: Shape of the working tensor to partition in the remaining
            shard space.
        dtype_size: Size in bytes of the dtype of the working tensor.
        num_elems: Number of elements in the working tensor.
  
      Returns:
        A tuple containing the axis of the next partition and that partition size.
      """
      if shape.rank is None or shape.rank == 0:
          return 0, math.inf
  
      # Find axis with minimum partitions. (aka axis with maximum partition size)
      # (max partition size is as close as possible to the shard_size_remaining)
      bytes_per_slice = num_elems // shape.dims[0].value * dtype_size
      slices_per_shard = max(
          1, math.floor(shard_size_remaining / bytes_per_slice)
      )
      min_parts = math.ceil(shape.dims[0].value / slices_per_shard)
      min_axis = 0
      for axis in range(1, shape.rank):
          bytes_per_slice = (
              num_elems // shape.dims[axis].value * dtype_size
          )
          slices_per_shard = max(
              1, math.floor(shard_size_remaining / bytes_per_slice)
          )
          axis_parts = math.ceil(
              shape.dims[axis].value / slices_per_shard
          )
          partition_size = num_elems * dtype_size / axis_parts
          if (
              axis_parts < min_parts
              and partition_size < shard_size_remaining
          ):
              min_axis, min_parts = axis, int(axis_parts)
      return min_axis, math.ceil(int(shape[min_axis]) / min_parts)
  
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 3 for calling the MaxShardSizePolicy obj's underlying member method code with example usage: obj._get_next_partition(shard_size_remaining, shape, dtype_size, num_elems) and returns: min_axis,  math, int,  math