message: hlo_runner to support device assignment with computation_count > 1.
diff: @@ -206,6 +206,7 @@ StatusOr<std::vector<Literal>> HloRunner::ExecuteReplicatedImpl(
  DeviceAssignment* device_assignment) {
  std::vector<std::unique_ptr<se::Stream>> streams;
  std::vector<ServiceExecutableRunOptions> service_run_options;
  + int64 num_partitions = device_assignment->computation_count();
  std::vector<ScopedShapedBuffer> argument_buffers;
  // This reserve() call is necessary for correctness, because
  @@ -228,7 +229,7 @@ StatusOr<std::vector<Literal>> HloRunner::ExecuteReplicatedImpl(
  int64 index = 0;
  RunId run_id;
  for (int64 i = 0; i < options.num_replicas; ++i) {
  - int64 device = (*device_assignment)(i, 0);
  + int64 device = (*device_assignment)(i / num_partitions, i % num_partitions);
  TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
  backend().stream_executor(device));
  streams.push_back(absl::make_unique<se::Stream>(executor));
  @@ -269,7 +270,8 @@ StatusOr<std::vector<Literal>> HloRunner::ExecuteReplicatedImpl(
  }
  if (!options.infeed_values.empty()) {
  for (int64 i = 0; i < options.num_replicas; ++i) {
  - int64 device = (*device_assignment)(i, 0);
  + int64 device =
  + (*device_assignment)(i / num_partitions, i % num_partitions);
  pool->Schedule([this, device, &options, i]() {
  se::StreamExecutor* executor =
  backend().stream_executor(device).ValueOrDie();
  @@ -290,7 +292,8 @@ StatusOr<std::vector<Literal>> HloRunner::ExecuteReplicatedImpl(
  options.outfeed_values->resize(options.num_replicas);
  }
  for (int64 i = 0; i < options.num_replicas; ++i) {
  - int64 device = (*device_assignment)(i, 0);
  + int64 device =
  + (*device_assignment)(i / num_partitions, i % num_partitions);
  pool->Schedule([this, device, &options, i]() {
  se::StreamExecutor* executor =
  backend().stream_executor(device).ValueOrDie();
  