Returns the serialized form of the topology.
(self)
| 208 | return self._device_coordinates.shape[1] |
| 209 | |
| 210 | def serialized(self): |
| 211 | """Returns the serialized form of the topology.""" |
| 212 | if self._serialized is None: |
| 213 | proto = topology_pb2.TopologyProto() |
| 214 | proto.mesh_shape[:] = list(self._mesh_shape) |
| 215 | proto.num_tasks = self._device_coordinates.shape[0] |
| 216 | proto.num_tpu_devices_per_task = self._device_coordinates.shape[1] |
| 217 | proto.device_coordinates.extend(list(self._device_coordinates.flatten())) |
| 218 | self._serialized = proto.SerializeToString() |
| 219 | |
| 220 | return self._serialized |