num_tasks));

  if (topology_proto.num_tpu_devices_per_task() != num_tpus_per_task)
    return absl::InvalidArgumentError(absl::StrCat(
        "number of TPU devices available per task must be "
        "'num_tpu_devices_per_task' in '",
        kTopologyAttr, "' (", topology_proto.num_tpu_devices_per_task(),
        "), got ", num_tpus_per_task));

//   `num_cores_per_replica` = 2
//   `topology_attr` (in proto debug string format) =
//     {
//       mesh_shape: 2
//       mesh_shape: 2
//       mesh_shape: 2
//       num_tasks: 2
//       num_tpu_devices_per_task: 4
//       device_coordinates: 0
//       device_coordinates: 0
//       device_coordinates: 0
//       device_coordinates: 0
//       device_coordinates: 1
//       device_coordinates: 0

                                          int num_tasks,
                                          int num_tpu_devices_per_task) {
  tpu::TopologyProto topology_proto;
  for (int mesh_dim : mesh_shape) topology_proto.add_mesh_shape(mesh_dim);
  topology_proto.set_num_tasks(num_tasks);
  topology_proto.set_num_tpu_devices_per_task(num_tpu_devices_per_task);
  return topology_proto.SerializeAsString();
}

std::string TopologyWithDeviceCoordinates(

// -----

// The following topology is used in subsequent test cases:
// Proto debug string:
//   mesh_shape: 1
//   mesh_shape: 1
//   mesh_shape: 1
//   mesh_shape: 2
//   num_tasks: 1
//   num_tpu_devices_per_task: 2
//   device_coordinates: 0
//   device_coordinates: 0
//   device_coordinates: 0
//   device_coordinates: 0
//   device_coordinates: 0
//   device_coordinates: 0
//   device_coordinates: 0