auto e = ops::Add(root.WithOpName("E"), d, c);
  auto f = ops::Neg(root.WithOpName("F"), e);
  auto g = ops::AddN(root.WithOpName("G"), std::initializer_list<Output>{e, f});

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(root.ToGraph(graph.get()));

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), /*arg_shapes=*/{},
                           /*fnlib_def=*/nullptr, &shape_info));

      tensorflow::OpRegistry::Global(), flib);

  tensorflow::XlaCompiler::Options options;
  options.device_type = tensorflow::DeviceType(kArbitraryDeviceName);
  options.client = client;
  options.flib_def = flib_def.get();
  tensorflow::XlaCompiler compiler(options);

  std::unique_ptr<tensorflow::Graph> graph_copy(
      new tensorflow::Graph(tensorflow::OpRegistry::Global()));

  auto consumer2 = ops::Identity(scope.WithOpName("consumer2"), out2);
  auto consumer3 = ops::Identity(scope.WithOpName("consumer3"), out3);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(scope.ToGraph(graph.get()));
  return graph;
}

// Makes an encapsulate body graph for use in tests.
static std::unique_ptr<Graph> MakeBodyGraph() {

namespace {

Status ForceXlaConstantsOnHost(const Scope& s,
                               FunctionLibraryDefinition* flib_def,
                               std::unique_ptr<Graph>* result) {
  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  GraphOptimizationPassOptions options;
  SessionOptions session_options;
  session_options.env = Env::Default();
  options.graph = &graph;
  options.session_options = &session_options;

}

class MlirGraphOptimizationPassTest : public Test {
 public:
  void Init(Status pass_run_result,
            const std::vector<MlirOptimizationPassState>& pass_states) {
    graph_ = std::make_unique<Graph>(OpRegistry::Global());

    int pass_priority = 0;
    for (const MlirOptimizationPassState& pass_state : pass_states) {
      auto optimization_pass =
          std::make_unique<NiceMock<MockMlirOptimizationPass>>();

  node_builder.Input(values);
  return opts.FinalizeBuilder(&node_builder);
}

TEST(XlaCompilationTest, StagePipelinePreserved) {
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  {
    // Graph:
    //       b
    //       |
    //       v
    // a -> add0 (ClusterX) -> relu0 (ClusterX) -> stage
    //
    //             b
    //             |
    //             v

  add2.node()->AddAttr("_cluster", "cluster2");
  auto out = ops::Mul(root.WithOpName("mul"), add1, add2);

  Graph graph_before_encapsulation(OpRegistry::Global());
  TF_ASSERT_OK(root.ToGraph(&graph_before_encapsulation));

  FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
  std::unique_ptr<Graph> graph;
  TF_ASSERT_OK(EncapsulateSubgraphsInFunctions(
      "_cluster", graph_before_encapsulation,

    const Scope& scope, FunctionLibraryDefinition* flib_def = nullptr) {
  FunctionDefLibrary flib;
  FunctionLibraryDefinition flib_def_local(OpRegistry::Global(), flib);
  if (flib_def == nullptr) {
    flib_def = &flib_def_local;
  }

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));

  TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));

  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(

LogicalResult ConvertTFTextAPI(mlir::func::FuncOp func, llvm::StringRef api,
                               mlir::TF::FuncAttr attr);

// Check if TF.Text Tensorflow ops are registered.
bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery);

}  // end namespace TFL
}  // end namespace mlir

// A helper object to create GraphOptimizationPassOptions.
struct GraphOptimizationPassWrapper {
  explicit GraphOptimizationPassWrapper()
      : library(OpRegistry::Global(), FunctionDefLibrary()) {
    session_options.env = Env::Default();
  }

  // Create GraphOptimizationPassOptions with a graph passed in constructor and
  // sensible options.