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_;
  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
};

    const NameAttrList& function,
    const std::vector<XlaCompiler::Argument>& args,
    DeviceCompileMode compile_mode, bool has_ref_vars,
    bool may_alias_resource_update, FunctionLibraryRuntime* flr,
    ResourceMgr* rm, const XlaCompiler::CompilationResult** compilation_result,
    xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable) {
  PjRtDeviceCompiler* pjrt_device_compiler;

Status GetNodesRelatedToRefVariablesInDirection(
    const Graph& graph, FunctionLibraryRuntime* lib_runtime,
    Direction direction, int depth, absl::flat_hash_set<Node*>* result);

absl::StatusOr<bool> DoesAnyCalleeHaveRefNodes(
    const CallTargetListTy& call_target_list,
    FunctionLibraryRuntime* lib_runtime, Direction direction, int depth) {
  const int kMaxDepth = 10;

  OptimizerOptions opts;
  auto pflr = std::make_unique<ProcessFunctionLibraryRuntime>(
      nullptr, options.session_options->env, /*config=*/nullptr,
      TF_GRAPH_DEF_VERSION, options.flib_def, opts);
  FunctionLibraryRuntime* flr =
      pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);

  for (Node* node : graph->nodes()) {
    if (CanCreateXlaKernel(node->def())) {
      const FunctionBody* fbody = nullptr;

  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");
  }

  TF_ASSIGN_OR_RETURN(
      std::vector<XlaCompiler::Argument> args,

// all of the nodes that have ref variables as input or output.
absl::StatusOr<absl::flat_hash_set<Node*>> GetNodesRelatedToRefVariables(
    const Graph& graph, FunctionLibraryRuntime* lib_runtime);

// Deterministically serialized the graph to a byte string.
absl::StatusOr<std::string> SerializeGraphDeterministic(const Graph& graph);

        TF_GRAPH_DEF_VERSION, lib_def_.get(), opts,
        /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
  }

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

  std::unique_ptr<OpKernel> kernel_;
};

// It takes vectors of constant and resource arguments explicitly.
// It does not have corresponding OpDef because it is never present
// in the GraphDef.
// Currently, it is used by eager runtime. FunctionLibraryRuntime creates
// this kernel when asked to create a kernel for an XLA-compiled function.
//
// `has_ref_vars`: whether the input computation can have reference variables.

    const string& host_graph_func_name,
    const std::map<string, int>& host_compute_core, FunctionLibraryRuntime* flr,
    FunctionLibraryDefinition* fld, std::vector<string>* shape_inference_graphs,
    bool* has_outside_compilation) {
  // Convert the function to graph.
  const string& func_name = func_name_attrs.name();
  FunctionLibraryRuntime::Handle handle;
  TF_RETURN_IF_ERROR(

  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(
      new ProcessFunctionLibraryRuntime(
          nullptr, Env::Default(), /*config=*/nullptr, TF_GRAPH_DEF_VERSION,
          flib_def, OptimizerOptions{}));
  FunctionLibraryRuntime* lib_runtime =
      pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice);

  TF_ASSIGN_OR_RETURN(absl::flat_hash_set<Node*> nodes_related_to_ref_vars,