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.
| 41 | |
| 42 | |
| 43 | class Topology(object): |
| 44 | """Describes a set of TPU devices. |
| 45 | |
| 46 | Represents both the shape of the physical mesh, and the mapping between |
| 47 | TensorFlow TPU devices to physical mesh coordinates. |
| 48 | """ |
| 49 | |
| 50 | def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None): |
| 51 | """Builds a Topology object. |
| 52 | |
| 53 | If `serialized` is not `None`, the topology is parsed from `serialized` and |
| 54 | the other arguments are ignored. Otherwise, the topology is computed from |
| 55 | `mesh_shape` and `device_coordinates`. |
| 56 | |
| 57 | Args: |
| 58 | serialized: A serialized `TopologyProto`, or `None`. If not `None`, the |
| 59 | serialized proto is parsed to discover the topology. |
| 60 | mesh_shape: A sequence of 3 positive integers, or `None`. If not `None`, |
| 61 | the shape of the TPU topology, in number of cores. Ignored if |
| 62 | `serialized` is not `None`. |
| 63 | device_coordinates: A rank 3 numpy array that describes the mapping from |
| 64 | TensorFlow TPU devices to TPU fabric coordinates, or `None`. Ignored |
| 65 | if `serialized is not `None`. |
| 66 | |
| 67 | Raises: |
| 68 | ValueError: If `serialized` does not describe a well-formed topology. |
| 69 | ValueError: If `serialized` is `None` and `mesh_shape` is not a sequence |
| 70 | of 3 positive integers. |
| 71 | ValueError: If `serialized` is `None` and `device_coordinates` is not a |
| 72 | rank 3 numpy int32 array that describes a valid coordinate mapping. |
| 73 | """ |
| 74 | |
| 75 | self._serialized = serialized |
| 76 | |
| 77 | if serialized: |
| 78 | self._parse_topology(serialized) |
| 79 | else: |
| 80 | self._mesh_shape = np.asarray(mesh_shape, dtype=np.int32) |
| 81 | self._device_coordinates = np.asarray(device_coordinates, np.int32) |
| 82 | if len(self._mesh_shape) != 3 or any(self._mesh_shape < 1): |
| 83 | raise ValueError("`mesh_shape` must be a sequence of 3 positive " |
| 84 | "entries; got {}".format(self._mesh_shape)) |
| 85 | |
| 86 | if (len(self._device_coordinates.shape) != 3 or |
| 87 | self._device_coordinates.shape[2] != len(self._mesh_shape)): |
| 88 | raise ValueError("`device_coordinates` must be a rank 3 int32 array " |
| 89 | "with minor dimension equal to the mesh shape rank") |
| 90 | |
| 91 | self._topology_tasks, self._topology_devices = self._invert_topology() |
| 92 | |
| 93 | def _parse_topology(self, serialized): |
| 94 | """Parses a serialized `TopologyProto` into `self`.""" |
| 95 | proto = topology_pb2.TopologyProto() |
| 96 | proto.ParseFromString(serialized) |
| 97 | |
| 98 | self._mesh_shape = np.array(proto.mesh_shape, dtype=np.int32) |
| 99 | if len(self._mesh_shape) != 3 or any(self._mesh_shape < 1): |
| 100 | raise ValueError("`mesh_shape` must be a vector of size 3 with positive " |