Parses a serialized `TopologyProto` into `self`.
(self, serialized)
| 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 " |
| 101 | "entries; got {}".format(self._mesh_shape)) |
| 102 | |
| 103 | if proto.num_tasks < 0: |
| 104 | raise ValueError("`num_tasks` must be >= 0; got {}".format( |
| 105 | proto.num_tasks)) |
| 106 | if proto.num_tpu_devices_per_task < 0: |
| 107 | raise ValueError("`num_tpu_devices_per_task` must be >= 0; got {}".format( |
| 108 | proto.num_tpu_devices_per_task)) |
| 109 | |
| 110 | expected_coordinates_size = ( |
| 111 | proto.num_tasks * proto.num_tpu_devices_per_task * len( |
| 112 | proto.mesh_shape)) |
| 113 | if len(proto.device_coordinates) != expected_coordinates_size: |
| 114 | raise ValueError("`device_coordinates` must have shape num_tasks ({}) * " |
| 115 | "num_tpu_devices_per_task ({}) * len(mesh_shape) ({}); " |
| 116 | "got shape {}".format(proto.num_tasks, |
| 117 | proto.num_tpu_devices_per_task, |
| 118 | proto.mesh_shape, |
| 119 | len(proto.device_coordinates))) |
| 120 | |
| 121 | coords = np.array(proto.device_coordinates, dtype=np.int32) |
| 122 | if any(coords < 0): |
| 123 | raise ValueError("`device_coordinates` must be >= 0") |
| 124 | coords = coords.reshape((proto.num_tasks, proto.num_tpu_devices_per_task, |
| 125 | len(proto.mesh_shape))) |
| 126 | self._device_coordinates = coords |
| 127 | |
| 128 | def _invert_topology(self): |
| 129 | """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps.""" |
no test coverage detected