tf.tpu.experimental.Topology

Describes a set of TPU devices.

Represents both the shape of the physical mesh, and the mapping between TensorFlow TPU devices to physical mesh coordinates.

Args
serialized A serialized TopologyProto, or None. If not None, the serialized proto is parsed to discover the topology.
mesh_shape A sequence of 4 positive integers, or None. If not None, the shape of the TPU topology, in number of cores. Ignored if serialized is not None.
device_coordinates A rank 4 numpy array that describes the mapping from TensorFlow TPU devices to TPU fabric coordinates, or None. Ignored if serialized is notNone`.
Raises
ValueError If serialized does not describe a well-formed topology.
ValueError If serialized is None and mesh_shape is not a sequence of 4 positive integers.
ValueError If serialized is None and device_coordinates is not a rank 4 numpy int32 array that describes a valid coordinate mapping.
Attributes
device_coordinates Describes the mapping from TPU devices to topology coordinates.
mesh_rank Returns the number of dimensions in the mesh.
mesh_shape A rank 1 int32 array describing the shape of the TPU topology.
missing_devices Array of indices of missing devices.
num_tasks Returns the number of TensorFlow tasks in the TPU slice.
num_tpus_per_task Returns the number of TPU devices per task in the TPU slice.

Methods

cpu_device_name_at_coordinates

View source

Returns the CPU device attached to a logical core.

serialized

View source

Returns the serialized form of the topology.

task_ordinal_at_coordinates

View source

Returns the TensorFlow task number attached to device_coordinates.

Args
device_coordinates An integer sequence describing a device's physical coordinates in the TPU fabric.
Returns
Returns the TensorFlow task number that contains the TPU device with those physical coordinates.

tpu_device_name_at_coordinates

View source

Returns the name of the TPU device assigned to a logical core.

tpu_device_ordinal_at_coordinates

View source

Returns the TensorFlow device number at device_coordinates.

Args
device_coordinates An integer sequence describing a device's physical coordinates in the TPU fabric.
Returns
Returns the TensorFlow device number within the task corresponding to attached to the device with those physical coordinates.

© 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/tpu/experimental/Topology