tf.contrib.tpu.device_assignment

Computes a device_assignment of a computation across a TPU topology.

tf.contrib.tpu.device_assignment(
    topology, computation_shape=None, computation_stride=None, num_replicas=1
)

Attempts to choose a compact grid of cores for locality.

Returns a DeviceAssignment that describes the cores in the topology assigned to each core of each replica.

computation_shape and computation_stride values should be powers of 2 for optimal packing.

Args
topology	A Topology object that describes the TPU cluster topology. To obtain a TPU topology, evaluate the Tensor returned by initialize_system using Session.run. Either a serialized TopologyProto or a Topology object may be passed. Note: you must evaluate the Tensor first; you cannot pass an unevaluated Tensor here.
computation_shape	A rank 1 int32 numpy array with size equal to the topology rank, describing the shape of the computation's block of cores. If None, the computation_shape is [1] * topology_rank.
computation_stride	A rank 1 int32 numpy array of size topology_rank, describing the inter-core spacing of the computation_shape cores in the TPU topology. If None, the computation_stride is [1] * topology_rank.
num_replicas	The number of computation replicas to run. The replicas will be packed into the free spaces of the topology.

Returns
A DeviceAssignment object, which describes the mapping between the logical cores in each computation replica and the physical cores in the TPU topology.

Raises
ValueError	If topology is not a valid Topology object.
ValueError	If computation_shape or computation_stride are not 1D int32 numpy arrays with shape [3] where all values are positive.
ValueError	If computation's replicas cannot fit into the TPU topology.