#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session_options.h"

namespace tensorflow {
namespace {

Status ClusterScoping(std::unique_ptr<Graph>* graph) {
  FixupSourceAndSinkEdges(graph->get());

  GraphOptimizationPassWrapper wrapper;
  wrapper.session_options.config.mutable_graph_options()
      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);

                        .Attr("T", result_types[i])
                        .Attr("index", i)
                        .Finalize(graph.get(), &node);
    TF_RETURN_IF_ERROR(status);
  }
  FixupSourceAndSinkEdges(graph.get());
  return graph;
}

bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type) {
  const auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;