Partitioned Ps Strategy

Partitioned PS StrategyBuilder with Greedy Load Balancer.

class PartitionedPS(local_proxy_variable=False, sync=True, staleness=0)[source]

Bases: autodist.strategy.base.StrategyBuilder

Partitioned PS StrategyBuilder with Greedy Load Balancer.

Determine the number of partitions for each partition-able variable by finding its minimum divisor along the first axis.

Then, use a greedy load balancer (determined by memory usage) to assign the partitions to Parameter Servers. This means that, unlike the standard PS StrategyBuilder, that a variable can be spread across multiple servers using this StrategyBuilder.

build(graph_item, resource_spec)[source]

Generate the Strategy.

static get_num_shards(var)[source]

Gets the minimum number of shards for a variable.


Load function that computes the byte size of a single-output Operation.

Copied (with modifications) from

This is intended to be used with “Variable” ops, which have a single Tensor output with the contents of the variable. However, it can also be used for calculating the size of any op that has a single output.

Intended to be used with GreedyLoadBalancingStrategy.


op – An Operation with a single output, typically a “Variable” op.


The number of bytes in the output Tensor.


ValueError – if op does not have a single output, or if the shape of the single output is not fully-defined.