DLDevice GetDlContext(TFE_TensorHandle* h, TF_Status* status) {
  DLDevice ctx;
  const char* device_name =
      tensorflow::unwrap(h)->BackingDeviceName(&status->status);
  DeviceNameUtils::ParsedName parsed_name;
  tensorflow::DeviceNameUtils::ParseFullName(device_name, &parsed_name);
  std::string device_type = parsed_name.type;
  int device_id = 0;
  if (parsed_name.has_id) {

}

tensorflow::Var* GetVariableFromSession(mlir::TF::VarHandleOp var_handle_op,
                                        llvm::StringRef device_name,
                                        const tensorflow::DeviceMgr* mgr) {
  tensorflow::Device* device = nullptr;
  if (!mgr || !mgr->LookupDevice(StringRefToView(device_name), &device).ok())
    return nullptr;
  tensorflow::Var* var_ptr = nullptr;
  const auto& container = var_handle_op.getContainer().str();

                                                 const string& device_name,
                                                 int device_ordinal) {
  return Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
      DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
      absl::StrCat("device: ", device_name, " device"));
}

}  // namespace

    OwnedTFStatus c_status(TF_NewStatus());
    SP_Stream stream_handle = static_cast<CStream*>(stream)->Handle();
    SP_DeviceMemoryBase device_mem = DeviceMemoryBaseToC(location);
    stream_executor_->mem_zero(&device_, stream_handle, &device_mem, size,
                               c_status.get());
    return StatusFromTF_Status(c_status.get());
  }

}

TF_StringView TF_GetDeviceName(TF_OpKernelContext* ctx) {
  const auto& device_name =
      reinterpret_cast<tensorflow::OpKernelContext*>(ctx)->device()->name();
  TF_StringView device_name_sv;
  device_name_sv.data = device_name.data();
  device_name_sv.len = device_name.length();
  return device_name_sv;
}

#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)

class DeviceSetup {
 public:
  void AddDevicesAndSetUp(
      const std::vector<std::string>& device_names,
      const std::optional<FunctionDef>& fdef = std::nullopt);
  Device* GetDevice(const string& device_name);
  FunctionLibraryRuntime* flr() { return flr_; }

 private:
  FunctionLibraryRuntime* flr_;
  std::unique_ptr<DeviceMgr> device_mgr_;
  std::unique_ptr<FunctionLibraryDefinition> lib_def_;

  NameAttrList function;
  function.set_name(std::string{func_name});

  std::string device_name;
  if (!platform_name.empty()) {
    device_name = absl::StrCat("/device:", platform_name, ":0");
  }
  FunctionLibraryRuntime* flr = pflr->GetFLR(device_name);
  if (flr == nullptr) {
    // Use CPU as the fallback to get the `FunctionLibraryRuntime`.
    flr = pflr->GetFLR("/device:CPU:0");

  }

  using DeviceTypeConstRef = std::reference_wrapper<const DeviceType>;

  absl::StatusOr<DeviceTypeConstRef> GetDeviceTypeFor(
      absl::string_view device_name) {
    TF_ASSIGN_OR_RETURN(DeviceId device_id, GetIdFor(device_name));
    return std::cref(*id_to_device_type_[device_id.id()]);
  }

  string DebugString(const DeviceSet& device_set) const;

 private:

  const std::string attrKey = "T";
  for (const Node* node : graph.nodes()) {
    if (node->type_string() == nodeType) {
      auto device_name = node->assigned_device_name();
      DeviceNameUtils::ParsedName device;
      if (DeviceNameUtils::ParseFullName(device_name, &device) &&
          device.has_job && device.job == jobType) {
        for (const auto& attr : node->attrs()) {
          auto attr_key = attr.first;

    return absl::OkStatus();
  }

  int device_count = platform->VisibleDeviceCount();

  for (int i = 0; i < device_count; ++i) {
    const string device_name = absl::StrCat("/physical_device:TPU:", i);
    devices->push_back(device_name);
  }

  return absl::OkStatus();
}

Status TpuNodeDeviceFactory::CreateDevices(
    const SessionOptions& session_options, const string& name_prefix,