parameters: (
      max_shard_bytes, axis=0, bytes_per_string_element=16, max_shards=None)
return_statement: return _partitioner
docstring: Get a partitioner for VariableScope to keep shards below `max_shard_bytes`.
  
    This partitioner will shard a Variable along one axis, attempting to keep
    the maximum shard size below `max_shard_bytes`.  In practice, this is not
    always possible when sharding along only one axis.  When this happens,
    this axis is sharded as much as possible (i.e., every dimension becomes
    a separate shard).
  
    If the partitioner hits the `max_shards` limit, then each shard may end up
    larger than `max_shard_bytes`. By default `max_shards` equals `None` and no
    limit on the number of shards is enforced.
  
    One reasonable value for `max_shard_bytes` is `(64 << 20) - 1`, or almost
    `64MB`, to keep below the protobuf byte limit.
  
    Args:
      max_shard_bytes: The maximum size any given shard is allowed to be.
      axis: The axis to partition along.  Default: outermost axis.
      bytes_per_string_element: If the `Variable` is of type string, this provides
        an estimate of how large each scalar in the `Variable` is.
      max_shards: The maximum number of shards in int created taking precedence
        over `max_shard_bytes`.
  
    Returns:
      A partition function usable as the `partitioner` argument to
      `variable_scope` and `get_variable`.
  
    Raises:
      ValueError: If any of the byte counts are non-positive.
docstring_summary: Get a partitioner for VariableScope to keep shards below `max_shard_bytes`.
docstring_tokens: Get a partitioner for VariableScope to keep shards below max_shard_bytes .
function: def variable_axis_size_partitioner(
      max_shard_bytes, axis=0, bytes_per_string_element=16, max_shards=None):
    """Get a partitioner for VariableScope to keep shards below `max_shard_bytes`.
  
    This partitioner will shard a Variable along one axis, attempting to keep
    the maximum shard size below `max_shard_bytes`.  In practice, this is not
    always possible when sharding along only one axis.  When this happens,
    this axis is sharded as much as possible (i.e., every dimension becomes
    a separate shard).
  
    If the partitioner hits the `max_shards` limit, then each shard may end up
    larger than `max_shard_bytes`. By default `max_shards` equals `None` and no
    limit on the number of shards is enforced.
  
    One reasonable value for `max_shard_bytes` is `(64 << 20) - 1`, or almost
    `64MB`, to keep below the protobuf byte limit.
  
    Args:
      max_shard_bytes: The maximum size any given shard is allowed to be.
      axis: The axis to partition along.  Default: outermost axis.
      bytes_per_string_element: If the `Variable` is of type string, this provides
        an estimate of how large each scalar in the `Variable` is.
      max_shards: The maximum number of shards in int created taking precedence
        over `max_shard_bytes`.
  
    Returns:
      A partition function usable as the `partitioner` argument to
      `variable_scope` and `get_variable`.
  
    Raises:
      ValueError: If any of the byte counts are non-positive.
    """
    if max_shard_bytes < 1 or bytes_per_string_element < 1:
      raise ValueError(
          "Both max_shard_bytes and bytes_per_string_element must be positive.")
    if max_shards and max_shards < 1:
      raise ValueError(
          "max_shards must be positive.")
  
    def _partitioner(shape, dtype):
      """Partitioner that partitions shards to have max_shard_bytes total size.
  
      Args:
        shape: A `TensorShape`.
        dtype: A `DType`.
  
      Returns:
        A tuple representing how much to slice each axis in shape.
  
      Raises:
        ValueError: If shape is not a fully defined `TensorShape` or dtype is not
          a `DType`.
      """
      if not isinstance(shape, tensor_shape.TensorShape):
        raise ValueError("shape is not a TensorShape: %s" % shape)
      if not shape.is_fully_defined():
        raise ValueError("shape is not fully defined: %s" % shape)
      if not isinstance(dtype, dtypes.DType):
        raise ValueError("dtype is not a DType: %s" % dtype)
  
      if dtype.base_dtype == dtypes.string:
        element_size = bytes_per_string_element
      else:
        element_size = dtype.size
  
      partitions = [1] * shape.ndims
      bytes_per_slice = 1.0 * (
          shape.num_elements() / shape.dims[axis].value) * element_size
      # How many slices can we fit on one shard of size at most max_shard_bytes?
      # At least one slice is required.
      slices_per_shard = max(1, math.floor(max_shard_bytes / bytes_per_slice))
      # How many shards do we need for axis given that each shard fits
      # slices_per_shard slices from a total of shape[axis] slices?
      axis_shards = int(math.ceil(
          1.0 * shape.dims[axis].value / slices_per_shard))
      if max_shards:
        axis_shards = min(max_shards, axis_shards)
  
      partitions[axis] = axis_shards
  
      return partitions
  
    return _partitioner
function_tokens: def variable_axis_size_partitioner ( max_shard_bytes , axis = 0 , bytes_per_string_element = 16 , max_shards = None ) : if max_shard_bytes < 1 or bytes_per_string_element < 1 : raise ValueError ( "Both max_shard_bytes and bytes_per_string_element must be positive." ) if max_shards and max_shards < 1 : raise ValueError ( "max_shards must be positive." ) def _partitioner ( shape , dtype ) : """Partitioner that partitions shards to have max_shard_bytes total size.
  
      Args:
        shape: A `TensorShape`.
        dtype: A `DType`.
  
      Returns:
        A tuple representing how much to slice each axis in shape.
  
      Raises:
        ValueError: If shape is not a fully defined `TensorShape` or dtype is not
          a `DType`.
      """ if not isinstance ( shape , tensor_shape . TensorShape ) : raise ValueError ( "shape is not a TensorShape: %s" % shape ) if not shape . is_fully_defined ( ) : raise ValueError ( "shape is not fully defined: %s" % shape ) if not isinstance ( dtype , dtypes . DType ) : raise ValueError ( "dtype is not a DType: %s" % dtype ) if dtype . base_dtype == dtypes . string : element_size = bytes_per_string_element else : element_size = dtype . size partitions = [ 1 ] * shape . ndims bytes_per_slice = 1.0 * ( shape . num_elements ( ) / shape . dims [ axis ] . value ) * element_size # How many slices can we fit on one shard of size at most max_shard_bytes? # At least one slice is required. slices_per_shard = max ( 1 , math . floor ( max_shard_bytes / bytes_per_slice ) ) # How many shards do we need for axis given that each shard fits # slices_per_shard slices from a total of shape[axis] slices? axis_shards = int ( math . ceil ( 1.0 * shape . dims [ axis ] . value / slices_per_shard ) ) if max_shards : axis_shards = min ( max_shards , axis_shards ) partitions [ axis ] = axis_shards return partitions return _partitioner