mesh_shape_dim.value()));

  if (topology_proto.num_tasks() != num_tasks)
    return absl::InvalidArgumentError(absl::StrCat(
        "number of tasks from available TPU devices must be 'num_tasks' in '",
        kTopologyAttr, "' (", topology_proto.num_tasks(), "), got ",
        num_tasks));

  if (topology_proto.num_tpu_devices_per_task() != num_tpus_per_task)