// We assume that any tensors already placed on this device are
      // ParallelTensors.
      typed_inputs.emplace_back(reinterpret_cast<ParallelTensor*>(
          TFE_TensorHandleDevicePointer(input, status)));
      if (TF_GetCode(status) != TF_OK) return;
    } else {
      typed_inputs.emplace_back(input);
    }
  }

  absl::optional<std::vector<MaybeParallelTensorOwned>> maybe_typed_outputs(

    std::vector<PartialTensorShape> shapes;
    shapes.reserve(num_values);
    for (int i = 0; i < num_values; ++i) {
      if (num_dims[i] < 0) {
        shapes.emplace_back();
      } else {
        shapes.emplace_back(ArraySlice<int64_t>(
            reinterpret_cast<const int64_t*>(dims[i]), num_dims[i]));
      }
    }
    op_->node_builder.Attr(attr_name, shapes);
    return absl::OkStatus();
  }

  inputs_.clear();
  for (const auto& p : inputs) {
    inputs_.emplace_back(TF_Output{p.first, 0});
    input_values_.emplace_back(p.second);
  }
}

void CSession::SetOutputs(std::initializer_list<TF_Operation*> outputs) {
  ResetOutputValues();
  outputs_.clear();
  for (TF_Operation* o : outputs) {
    outputs_.emplace_back(TF_Output{o, 0});
  }
  output_values_.resize(outputs_.size());
}

    return;
  }
  unwrapped_results.resize(real_num_outputs);
  outputs->reserve(real_num_outputs);
  for (TFE_TensorHandle* unwrapped_result : unwrapped_results) {
    outputs->emplace_back(unwrapped_result);
  }
}

ParallelDevice::ParallelDevice(const std::vector<std::string>& devices,
                               bool is_async, int in_flight_nodes_limit)
    : underlying_devices_(devices),

                                    "Encountered while processing input ", i,
                                    " into function '", fn_name, "'");

    input_tensors->emplace_back(node, idx);

    const auto& iter = input_nodes->find(node);
    if (iter == input_nodes->end()) {
      input_nodes->insert({node, {idx}});
    } else {
      auto& indices = iter->second;

  std::vector<TFE_TensorHandle*> tensors() const {
    std::vector<TFE_TensorHandle*> result;
    result.reserve(tensors_.size());
    for (const TensorHandlePtr& tensor : tensors_) {
      result.emplace_back(tensor.get());
    }
    return result;
  }

 private:
  ParallelTensor(const ParallelDevice& device,
                 std::vector<TensorHandlePtr> tensors,

    // make the pointers invalid.
    all_input_tensors.reserve(num_inputs);
    for (int i = 0; i < num_inputs; ++i) {
      if (input_tensors[i] == nullptr) continue;
      all_input_tensors.emplace_back();
      Tensor& input_tensor = all_input_tensors.back();
      status->status = TF_TensorToTensor(input_tensors[i], &input_tensor);
      if (!status->status.ok()) return;
      input_tensors_vector[i] = &input_tensor;

  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  const int num_devices = TF_DeviceListCount(devices);
  for (int i = 0; i < num_devices; ++i) {
    device_names.emplace_back(TF_DeviceListName(devices, i, status));
    EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  }
  TF_DeleteDeviceList(devices);
  TF_DeleteStatus(status);
  return device_names;