W3cubDocs

/TensorFlow 2.4

tf.distribute.experimental.partitioners.FixedShardsPartitioner

Partitioner that allocates a fixed number of shards.

Inherits From: Partitioner

Examples:

# standalone usage:
partitioner = FixedShardsPartitioner(num_shards=2)
partitions = partitioner(tf.TensorShape([10, 3]), tf.float32)
[2, 1]

# use in ParameterServerStrategy
# strategy = tf.distribute.experimental.ParameterServerStrategy(
#   cluster_resolver=cluster_resolver, variable_partitioner=partitioner)
Args
num_shards int, number of shards to partition.

Methods

__call__

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]
Args
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.
Returns
A list of integers representing the number of partitions on each axis, where i-th value correponds to i-th axis.

© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/distribute/experimental/partitioners/FixedShardsPartitioner