Google I/O returns May 18-20! Reserve space and build your schedule Register now


Partitioner that allocates a minimum size per shard.

Inherits From: Partitioner

This partitioner ensures each shard has at least min_shard_bytes, and tries to allocate as many shards as possible, i.e., keeping shard size as small as possible. The maximum number of such shards (upper bound) is given by max_shards.


partitioner = MinSizePartitioner(min_shard_bytes=4, max_shards=2)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[2, 1]
partitioner = MinSizePartitioner(min_shard_bytes=4, max_shards=10)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[6, 1]

# use in ParameterServerStrategy
# strategy = tf.distribute.experimental.ParameterServerStrategy(
#   cluster_resolver=cluster_resolver, variable_partitioner=partitioner)

min_shard_bytes Minimum bytes of each shard. Defaults to 256K.
max_shards Upper bound on the number of shards. Defaults to 1.
bytes_per_string If the partition value is of type string, this provides an estimate of how large each string is.



View source

Partitions the given shape and returns the partition results.

Examples of a partitioner that allocates a fixed number of shards:

partitioner = FixedShardsPartitioner(num_shards=2)
partitions = partitioner(tf.TensorShape([10, 3], tf.float32), axis=0)
print(partitions) # [2, 0]

shape a tf.TensorShape, the shape to partition.
dtype a tf.dtypes.Dtype indicating the type of the partition value.
axis The axis to partition along. Default: outermost axis.

A list of integers representing the number of partitions on each axis, where i-th value correponds to i-th axis.