auto ret0 = ops::_Retval(s.WithOpName("ret0"), arg1, 0);
    auto ret1 = ops::_Retval(s.WithOpName("ret1"), arg0, 1);
    std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
    TF_CHECK_OK(s.ToGraph(g.get()));
    FunctionDef *xla_fdef = fdl.add_function();
    TF_CHECK_OK(GraphToFunctionDef(*g, "f1", xla_fdef));
  }
  {
    // Function for While's "body".
    // "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)

  ops::PartitionedCall c(root.WithOpName("C"), {a}, {DT_FLOAT}, c_name_attr);
  Output d = ops::Add(root.WithOpName("D"), b.output.front(), c.output.front());

  TF_ASSERT_OK(root.ToGraph(graph.get()));
  TF_ASSERT_OK(
      MarkForCompilationPassTestHelper::MarkForCompilation(&graph, &flib_def));
  auto clusters = GetClusters(*graph);

  EXPECT_EQ(2, clusters.size());

  auto graph = std::make_unique<Graph>(OpRegistry::Global());
  TF_ASSERT_OK(s.ToGraph(graph.get()));

  // This is needed to register the XLA_GPU device.
  std::vector<std::unique_ptr<Device>> devices;
  TF_ASSERT_OK(DeviceFactory::AddDevices(
      SessionOptions(), "/job:localhost/replica:0/task:0", &devices));

  // Scope::ToGraph loses the assigned device name since it goes through

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

          .Attr("Tout", {DT_RESOURCE})
          .Attr("f", f_name_attr)
          .Finalize(root.graph(), &call));
  call->AddAttr(std::string(kMustCompileAttr), true);

  TF_ASSERT_OK(root.ToGraph(&graph));

  // Required for passing the PS server parameter check.
  for (Node* node : graph.nodes()) {
    node->set_assigned_device_name("/job:ps/replica:0/task:0/device:GPU:0");
  }

  EXPECT_TRUE(

  auto ret0 = ops::_Retval(s.WithOpName("ret0"), add, 0);
  auto ret1 = ops::_Retval(s.WithOpName("ret1"), arg1, 1);
  std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(s.ToGraph(g.get()));
  auto node_name_image = g->BuildNodeNameIndex();
  Node *add_node = node_name_image["add"];
  EXPECT_NE(add_node, nullptr);
  add_node->AddAttr(kXlaConnectedToXlaComputationAttrName, "cluster");

  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() {
  Scope scope = Scope::NewRootScope().ExitOnError();

  Output add = ops::Add(s.WithOpName("add"), const_0, const_1);
  Output identity = ops::Identity(s.WithOpName("identity"), add);
  Graph g(OpRegistry::Global());
  TF_CHECK_OK(s.ToGraph(&g));

  TF_CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g));

  // Check that "add" node now has _xla_inferred_shapes attr.
  auto node_index = g.BuildNodeNameIndex();
  Node *add_node = node_index["add"];

      ->mutable_optimizer_options()
      ->set_global_jit_level(OptimizerOptions::ON_2);
  GraphOptimizationPassOptions options;
  options.graph = &graph;
  options.session_options = &session_options;

  // Scope::ToGraph seems to drop assigned devices, probably because it goes
  // through a GraphDef.  So explicitly maintain the device assignment.
  // std::unordered_map<string, string> assigned_device_names;
  // for (Node* n : s.graph()->nodes()) {

  SessionOptions session_options;
  session_options.env = Env::Default();
  options.graph = &graph;
  options.session_options = &session_options;
  options.flib_def = flib_def;
  TF_RETURN_IF_ERROR(s.ToGraph(graph.get()));
  ForceXlaConstantsOnHostPass rewriter;
  TF_RETURN_IF_ERROR(rewriter.Run(options));
  *result = std::move(graph);
  return absl::OkStatus();
}

TEST(ForceXlaConstantsOnHostPassTest, Simple) {