/

• English
• 中文 – 简体

GitHub Sign in

• TensorFlow v2.16.1

tf.tpu.experimental.Topology

View source on GitHub

Describes a set of TPU devices.

tf.tpu.experimental.Topology(
    serialized=None, mesh_shape=None, device_coordinates=None
)

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 3 numpy array that describes the mapping from TensorFlow TPU devices to TPU fabric coordinates, or None. If specified, array is a rank 3 int32 array with shape [tasks, devices, axis]. tasks is the number of tasks in the TPU cluster, devices is the number of TPU devices per task, and axis is the number of axes in the TPU cluster topology. Each entry gives the axis-th coordinate in the topology of a task/device pair. TPU topologies are 4-dimensional, with dimensions (x, y, z, core number). This arg is 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 3 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

cpu_device_name_at_coordinates(
    device_coordinates, job=None
)

Returns the CPU device attached to a logical core.

serialized

View source

serialized()

Returns the serialized form of the topology.

task_ordinal_at_coordinates

View source

task_ordinal_at_coordinates(
    device_coordinates
)

Returns the TensorFlow task number attached to device_coordinates.

device_coordinates An integer sequence describing a device's physical coordinates in the TPU fabric.

tpu_device_name_at_coordinates

View source

tpu_device_name_at_coordinates(
    device_coordinates, job=None
)

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

tpu_device_ordinal_at_coordinates

View source

tpu_device_ordinal_at_coordinates(
    device_coordinates
)

Returns the TensorFlow device number at device_coordinates.

device_coordinates An integer sequence describing a device's physical coordinates in the TPU fabric.