tf.distribute.experimental.partitioners.MaxSizePartitioner

Partitioner that keeps shards below max_shard_bytes.

Inherits From: Partitioner

This partitioner ensures each shard has at most max_shard_bytes, and tries to allocate as few shards as possible, i.e., keeping shard size as large as possible.

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.

Examples:

partitioner = MaxSizePartitioner(max_shard_bytes=4)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[6, 1]
partitioner = MaxSizePartitioner(max_shard_bytes=4, max_shards=2)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[2, 1]
partitioner = MaxSizePartitioner(max_shard_bytes=1024)
partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
[1, 1]

# use in