xla::ExecutableRunOptions run_options;
  xla::gpu::GpuExecutableRunOptions gpu_options;
  xla::DeviceAssignment device_assignment;
  if (result->collective_info) {
    TF_RETURN_IF_ERROR(ResolveDeviceAssignment(ctx, *result->collective_info,
                                               run_options, device_assignment,
                                               gpu_options));
  }

  run_options.set_stream(stream);

  if (iter != session_options.config.device_count().end() &&
      iter->second == 0) {
    // Device count for GPU is 0.
    return absl::OkStatus();
  }

  string allowed_gpus =
      session_options.config.gpu_options().visible_device_list();
  std::optional<std::set<int>> gpu_ids =
      ParseVisibleDeviceList(allowed_gpus).value();
  if (!gpu_ids) {
    gpu_ids.emplace();

  if ((tensorflow::IsGoogleCudaEnabled() || tensorflow::IsBuiltWithROCm()) &&
      num_virtual_gpus > 0) {
    tensorflow::GPUOptions* gpu_options =
        server_def.mutable_default_session_config()->mutable_gpu_options();
    auto virtual_devices =
        gpu_options->mutable_experimental()->add_virtual_devices();
    for (int i = 0; i < num_virtual_gpus; ++i) {
      virtual_devices->add_memory_limit_mb(200);
    }

    flags->tf_xla_cpu_global_jit = true;
    flags->tf_xla_min_cluster_size = 1;
  } else {
    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF);
  }

  auto* gpu_options = config.mutable_gpu_options();
  gpu_options->set_allow_growth(gpu_memory_allow_growth);

  (*config.mutable_device_count())["CPU"] = num_cpu_devices;

  // TODO(b/113217601): This is needed for EagerContext::runner_ to use a

    FunctionLibraryRuntime* flr) {
  std::optional<std::set<int>> gpu_ids = std::nullopt;

  if (flr->config_proto()) {
    string allowed_gpus =
        flr->config_proto()->gpu_options().visible_device_list();
    TF_ASSIGN_OR_RETURN(gpu_ids, ParseVisibleDeviceList(allowed_gpus));
  }

  return gpu_ids;
}

Status GetCompilationDeviceTypeAndPjRtClient(

              /*missing_ctx_input_prefix=*/0, input_output_alias);
      OP_REQUIRES_OK_ASYNC(ctx, execution_inputs.status(), done);

      xla::gpu::GpuExecutableRunOptions gpu_options;
      xla::DeviceAssignment device_assignment;
      xla::ExecutableRunOptions run_options;
      if (compilation_result->collective_info.has_value()) {
        OP_REQUIRES_OK_ASYNC(ctx,

// Create a serialized tensorflow.ConfigProto proto, where:
//
// a) ConfigProto.optimizer_options.global_jit_level is set to ON_1 if
// `enable_xla_compilation` is non-zero, and OFF otherwise.
// b) ConfigProto.gpu_options.allow_growth is set to `gpu_memory_allow_growth`.
// c) ConfigProto.device_count is set to `num_cpu_devices`.
TF_CAPI_EXPORT extern TF_Buffer* TF_CreateConfig(

  GraphOptimizationPassOptions opt_options;
  opt_options.graph = graph;
  opt_options.session_options = &session_options;
  opt_options.flib_def = flib_def;

  if (options.enable_cluster_scoping) {
    ClusterScopingPass cluster_scoping_pass;
    TF_RETURN_IF_ERROR(cluster_scoping_pass.Run(opt_options));
  }

  MarkForCompilationPass mark_for_compilation_pass;

      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);
  GraphOptimizationPassOptions opt_options =
      wrapper.CreateGraphOptimizationPassOptions(graph);

  ClusterScopingPass pass;
  return pass.Run(opt_options);
}

absl::flat_hash_map<string, string> GetXlaInternalScopes(const Graph& graph) {
  absl::flat_hash_map<string, string> scopes;

  FixupSourceAndSinkEdges(graph.get());

  GraphOptimizationPassWrapper wrapper;
  GraphOptimizationPassOptions opt_options =
      wrapper.CreateGraphOptimizationPassOptions(&graph);
  opt_options.flib_def = &flib_def;

  BuildXlaOpsPass pass(/*enable_lazy_compilation=*/true);
  TF_RETURN_IF_ERROR(pass.Run(opt_options));
  VLOG(3) << graph->ToGraphDefDebug().DebugString();
  *result = std::move(graph);
  return absl::OkStatus();
}