TFE_OutputTensorHandles InputTFE_OutputTensorHandles(
    const py::handle& num_outputs) {
  TFE_OutputTensorHandles output_tensor_handles;
#if PY_MAJOR_VERSION < 3
  if (!PyInt_Check(num_outputs.ptr())) {
#else
  if (!PyLong_Check(num_outputs.ptr())) {
#endif
    PyErr_SetString(PyExc_TypeError,
                    "expected an integer value (size of the number of "
                    "outputs of the operation)");
    throw py::error_already_set();
  }
#if PY_MAJOR_VERSION < 3
  long sz = PyInt_AsLong(num_outputs.ptr());  // NOLINT
#else
  long sz = PyLong_AsLong(num_outputs.ptr());  // NOLINT
#endif
  if (sz > 0) {
#if PY_MAJOR_VERSION < 3
    output_tensor_handles.resize(PyInt_AsLong(num_outputs.ptr()), nullptr);
#else
    output_tensor_handles.resize(PyLong_AsLong(num_outputs.ptr()), nullptr);
#endif
  }
  return output_tensor_handles;
}

// Packs multiple `EagerTensor`s of the same dtype and shape into one
// `EagerTensor`.
py::object TFE_Py_PackEagerTensors_wrapper(const py::handle& context,
                                           const py::handle& tensors) {
  TFE_Context* ctx = tensorflow::InputTFE_Context(context);
  TFE_InputTensorHandles handles = InputTFE_InputTensorHandles(tensors);
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  int size = handles.size();
  TFE_TensorHandle* packed_handle =
      TFE_CreatePackedTensorHandle(ctx, handles.data(), &size, status.get());
  tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
  PyObject* packed_tensor =
      EagerTensorFromHandle(packed_handle, /*is_packed=*/true);
  return tensorflow::PyoOrThrow(packed_tensor);
}

// This function was created from fusing the typemap logic in platform/base.i.
py::object TFE_Py_ExecuteCancelable_wrapper(
    const py::handle& context, const char* device_name, const char* op_name,
    const py::handle& inputs, const py::handle& attrs,
    TFE_CancellationManager* cancellation_manager,
    const py::handle& num_outputs) {
  TFE_Context* ctx = tensorflow::InputTFE_Context(context);
  TFE_InputTensorHandles input_tensor_handles =
      InputTFE_InputTensorHandles(inputs);
  TFE_OutputTensorHandles output_tensor_handles =
      InputTFE_OutputTensorHandles(num_outputs);
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  TFE_Py_ExecuteCancelable(ctx, device_name, op_name, &input_tensor_handles,
                           attrs.ptr(), cancellation_manager,
                           &output_tensor_handles, status.get());

  int output_len = output_tensor_handles.size();
  PyObject* output_list = PyList_New(output_len);
  for (int i = 0; i < output_len; ++i) {
    PyObject* output;
    output = EagerTensorFromHandle(output_tensor_handles.at(i));
    PyList_SetItem(output_list, i, output);
  }
  tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
  return tensorflow::PyoOrThrow(output_list);
}

static py::object TF_ListPhysicalDevices() {
  std::vector<string> devices;
  tensorflow::Status s =
      tensorflow::DeviceFactory::ListAllPhysicalDevices(&devices);
  MaybeRaiseRegisteredFromStatus(s);
  PyObject* result = PyList_New(devices.size());
  int i = 0;
  for (auto& dev : devices) {
    PyObject* dev_obj = PyBytes_FromStringAndSize(dev.data(), dev.size());
    PyList_SetItem(result, i, dev_obj);
    ++i;
  }
  return tensorflow::PyoOrThrow(result);
}

static std::unordered_map<string, string> TF_GetDeviceDetails(int index) {
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  std::unordered_map<string, string> device_details;
  tensorflow::Status s =
      tensorflow::DeviceFactory::GetAnyDeviceDetails(index, &device_details);
  tensorflow::Set_TF_Status_from_Status(status.get(), s);
  MaybeRaiseRegisteredFromTFStatus(status.get());
  return device_details;
}

static py::object TFE_ClearScalarCache() {
  tensorflow::TFE_TensorHandleCache::Get()->Clear();
  return py::none();
}

// Returns compiler IR for a given function.
static std::string TFE_GetCompilerIr(py::handle& ctx,
                                     const char* concrete_function_name,
                                     const char* stage, const char* device_name,
                                     py::handle& inputs) {
  EagerContext* context = ContextFromInterface(
      reinterpret_cast<ImmediateExecutionContext*>(InputTFE_Context(ctx)));

  std::string s_stage(stage);
  IrExportStage selected_stage = [&] {
    if (s_stage == "hlo") {
      return IrExportStage::HLO;
    } else if (s_stage == "optimized_hlo") {
      return IrExportStage::OPTIMIZED_HLO;
    } else {
      ThrowValueError(
          absl::StrFormat("Invalid stage selected: '%s'. Valid values are: "
                          "'hlo', 'optimized_hlo'",
                          s_stage)
              .c_str());
    }
  }();

  TFE_InputTensorHandles handles = InputTFE_InputTensorHandles(inputs);

  std::vector<const Tensor*> input_tensors;
  for (TFE_TensorHandle* tensor_handle : handles) {
    AbstractTensorHandle* abstract_tensor_handle = unwrap(tensor_handle);
    TensorHandle* th = TensorHandleFromInterface(abstract_tensor_handle);

    const Tensor* t;
    Status st = th->Tensor(&t);
    if (!st.ok()) {
      ThrowValueError(
          absl::StrFormat("Could not resolve tensor: '%s'", st.error_message())
              .c_str());
    }
    input_tensors.push_back(t);
  }

  DeviceNameUtils::ParsedName input_device_name;
  if (!DeviceNameUtils::ParseFullOrLocalName(device_name, &input_device_name)) {
    ThrowValueError(
        absl::StrFormat("Failed parsing device name: '%s'", device_name)
            .c_str());
  }

  std::vector<Device*> devices = context->local_device_mgr()->ListDevices();
  auto selected_device = absl::c_find_if(devices, [&](const Device* d) {
    return DeviceNameUtils::AreCompatibleDevNames(input_device_name,
                                                  d->parsed_name());
  });
  if (selected_device == devices.end()) {
    ThrowValueError("No matching device found");
  }

  xla::StatusOr<std::string> hlo_text =
      GetCompilerIr(selected_stage, context->pflr(), concrete_function_name,
                    *selected_device, input_tensors);

  if (!hlo_text.ok()) {
    ThrowValueError(absl::StrFormat("Failed getting HLO text: '%s'",
                                    hlo_text.status().error_message())
                        .c_str());
  }
  return *hlo_text;
}

}  // namespace tensorflow
TFE_OutputTensorHandles InputTFE_OutputTensorHandles(
    const py::handle& num_outputs) {
  TFE_OutputTensorHandles output_tensor_handles;
#if PY_MAJOR_VERSION < 3
  if (!PyInt_Check(num_outputs.ptr())) {
#else
  if (!PyLong_Check(num_outputs.ptr())) {
#endif
    PyErr_SetString(PyExc_TypeError,
                    "expected an integer value (size of the number of "
                    "outputs of the operation)");
    throw py::error_already_set();
  }
#if PY_MAJOR_VERSION < 3
  long sz = PyInt_AsLong(num_outputs.ptr());   
#else
  long sz = PyLong_AsLong(num_outputs.ptr());   
#endif
  if (sz > 0) {
#if PY_MAJOR_VERSION < 3
    output_tensor_handles.resize(PyInt_AsLong(num_outputs.ptr()), nullptr);
#else
    output_tensor_handles.resize(PyLong_AsLong(num_outputs.ptr()), nullptr);
#endif
  }
  return output_tensor_handles;
}

 
 
py::object TFE_Py_PackEagerTensors_wrapper(const py::handle& context,
                                           const py::handle& tensors) {
  TFE_Context* ctx = tensorflow::InputTFE_Context(context);
  TFE_InputTensorHandles handles = InputTFE_InputTensorHandles(tensors);
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  int size = handles.size();
  TFE_TensorHandle* packed_handle =
      TFE_CreatePackedTensorHandle(ctx, handles.data(), &size, status.get());
  tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
  PyObject* packed_tensor =
      EagerTensorFromHandle(packed_handle,  true);
  return tensorflow::PyoOrThrow(packed_tensor);
}

 
py::object TFE_Py_ExecuteCancelable_wrapper(
    const py::handle& context, const char* device_name, const char* op_name,
    const py::handle& inputs, const py::handle& attrs,
    TFE_CancellationManager* cancellation_manager,
    const py::handle& num_outputs) {
  TFE_Context* ctx = tensorflow::InputTFE_Context(context);
  TFE_InputTensorHandles input_tensor_handles =
      InputTFE_InputTensorHandles(inputs);
  TFE_OutputTensorHandles output_tensor_handles =
      InputTFE_OutputTensorHandles(num_outputs);
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  TFE_Py_ExecuteCancelable(ctx, device_name, op_name, &input_tensor_handles,
                           attrs.ptr(), cancellation_manager,
                           &output_tensor_handles, status.get());

  int output_len = output_tensor_handles.size();
  PyObject* output_list = PyList_New(output_len);
  for (int i = 0; i < output_len; ++i) {
    PyObject* output;
    output = EagerTensorFromHandle(output_tensor_handles.at(i));
    PyList_SetItem(output_list, i, output);
  }
  tensorflow::MaybeRaiseRegisteredFromTFStatus(status.get());
  return tensorflow::PyoOrThrow(output_list);
}

static py::object TF_ListPhysicalDevices() {
  std::vector<string> devices;
  tensorflow::Status s =
      tensorflow::DeviceFactory::ListAllPhysicalDevices(&devices);
  MaybeRaiseRegisteredFromStatus(s);
  PyObject* result = PyList_New(devices.size());
  int i = 0;
  for (auto& dev : devices) {
    PyObject* dev_obj = PyBytes_FromStringAndSize(dev.data(), dev.size());
    PyList_SetItem(result, i, dev_obj);
    ++i;
  }
  return tensorflow::PyoOrThrow(result);
}

static std::unordered_map<string, string> TF_GetDeviceDetails(int index) {
  tensorflow::Safe_TF_StatusPtr status = tensorflow::make_safe(TF_NewStatus());
  std::unordered_map<string, string> device_details;
  tensorflow::Status s =
      tensorflow::DeviceFactory::GetAnyDeviceDetails(index, &device_details);
  tensorflow::Set_TF_Status_from_Status(status.get(), s);
  MaybeRaiseRegisteredFromTFStatus(status.get());
  return device_details;
}

static py::object TFE_ClearScalarCache() {
  tensorflow::TFE_TensorHandleCache::Get()->Clear();
  return py::none();
}

 
static std::string TFE_GetCompilerIr(py::handle& ctx,
                                     const char* concrete_function_name,
                                     const char* stage, const char* device_name,
                                     py::handle& inputs) {
  EagerContext* context = ContextFromInterface(
      reinterpret_cast<ImmediateExecutionContext*>(InputTFE_Context(ctx)));

  std::string s_stage(stage);
  IrExportStage selected_stage = [&] {
    if (s_stage == "hlo") {
      return IrExportStage::HLO;
    } else if (s_stage == "optimized_hlo") {
      return IrExportStage::OPTIMIZED_HLO;
    } else {
      ThrowValueError(
          absl::StrFormat("Invalid stage selected: '%s'. Valid values are: "
                          "'hlo', 'optimized_hlo'",
                          s_stage)
              .c_str());
    }
  }();

  TFE_InputTensorHandles handles = InputTFE_InputTensorHandles(inputs);

  std::vector<const Tensor*> input_tensors;
  for (TFE_TensorHandle* tensor_handle : handles) {
    AbstractTensorHandle* abstract_tensor_handle = unwrap(tensor_handle);
    TensorHandle* th = TensorHandleFromInterface(abstract_tensor_handle);

    const Tensor* t;
    Status st = th->Tensor(&t);
    if (!st.ok()) {
      ThrowValueError(
          absl::StrFormat("Could not resolve tensor: '%s'", st.error_message())
              .c_str());
    }
    input_tensors.push_back(t);
  }

  DeviceNameUtils::ParsedName input_device_name;
  if (!DeviceNameUtils::ParseFullOrLocalName(device_name, &input_device_name)) {
    ThrowValueError(
        absl::StrFormat("Failed parsing device name: '%s'", device_name)
            .c_str());
  }

  std::vector<Device*> devices = context->local_device_mgr()->ListDevices();
  auto selected_device = absl::c_find_if(devices, [&](const Device* d) {
    return DeviceNameUtils::AreCompatibleDevNames(input_device_name,
                                                  d->parsed_name());
  });
  if (selected_device == devices.end()) {
    ThrowValueError("No matching device found");
  }

  xla::StatusOr<std::string> hlo_text =
      GetCompilerIr(selected_stage, context->pflr(), concrete_function_name,
                    *selected_device, input_tensors);

  if (!hlo_text.ok()) {
    ThrowValueError(absl::StrFormat("Failed getting HLO text: '%s'",
                                    hlo_text.status().error_message())
                        .c_str());
  }
  return *hlo_text;
}

}   