FunctionDefLibrary flib;
  *flib.add_function() = FunctionDefHelper::Define(
      /*Function*/ kCompilableFunctionName,
      /*Inputs*/ {"n_a:float", "n_b:float"},
      /*Outputs*/ {"n_c:float"},
      /*Attribute*/ {},
      // Node info
      {{{kCompilableFunctionNodeName},
        "Add",
        {"n_a", "n_b"},
        {{"T", DT_FLOAT}}}});
  *flib.add_function() = FunctionDefHelper::Define(

    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)
    // "ret0" = "arg0"
    // "ret1" = "arg1"

      {{{"out"}, "Identity", {"in"}, {{"T", DataType::DT_RESOURCE}}}});
}

TEST(IsSupportedByNonReplicatedBridge, NonReplicatedGraph) {
  const FunctionDef& fd = PassThroughResource();
  FunctionDefLibrary flib;
  *flib.add_function() = fd;
  FunctionLibraryDefinition flib_def(OpRegistry::Global(), flib);
  Graph graph(flib_def);
  graph.SetConstructionContext(ConstructionContext::kEagerRuntime);
  tensorflow::set_tf2_execution(true);

    PartialTensorShape shape({2});
    node_name_image["identity1"]->AddAttr(
        kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});

    FunctionDef *xla_fdef = fdl.add_function();
    TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef));
  }
  FunctionLibraryDefinition fld(OpRegistry::Global(), fdl);

  protobuf::Map<string, tensorflow::AttrValue> attrs;

  GraphDef graphdef_expected;

  *library_expected.add_function() = test::function::XTimesTwo();
  *library_expected.add_function() = FunctionDefHelper::Create(
      "F1", {"a_0_arg:float"}, {"c_0_retval:float"}, {},
      {
          {{"C"}, "UnaryTest", {"a_0_arg"}},
      },
      {{"c_0_retval", "C:o:0"}});
  *library_expected.add_function() = FunctionDefHelper::Create(

    EXPECT_LT(max_error, 1e-3);
  }

  Scope scope_;
};

TEST_F(FunctionGradTest, PartitionedCallGrad) {
  FunctionDefLibrary f_lib_proto;
  *(f_lib_proto.add_function()) = test::function::XTimesTwo();

  // Construct a graph:
  //   A = Placeholder[dtype=int32]
  //   B = XTimesTwo[_tpu_replicate="cluster"](A)
  //   C = XTimesTwo[_xla_compile_id="cluster"](A)

                                {{"c", "t0:y:0"}});

  AttrValue true_attribute;
  true_attribute.set_b(true);
  (*called_func.mutable_attr())[kXlaMustCompileAttr] = true_attribute;
  *library.add_function() = called_func;
  TF_ASSERT_OK(root.graph()->AddFunctionLibrary(library));
  FunctionLibraryDefinition flib_def(OpRegistry::Global(), library);
  Output in = ops::Placeholder(root, DT_FLOAT);

       {{"k"}, "Add", {"j", "h"}, {{"T", DT_FLOAT}}},
       {{"l"}, "Add", {"k", "h"}, {{"T", DT_FLOAT}}},
       {{"m"}, "Identity", {"l"}, {{"T", DT_FLOAT}}}});

  GraphDef graph;
  *graph.mutable_library()->add_function() = make_test_fn;
  *graph.add_node() = MakeNode("a", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("b", "Placeholder", {},

    TF_CHECK_OK(DeviceFactory::AddDevices(
        options, "/job:localhost/replica:0/task:0", &devices));

    FunctionDefLibrary proto;
    for (const auto& fdef : flib) {
      *(proto.add_function()) = fdef;
    }
    lib_def_ = std::make_unique<FunctionLibraryDefinition>(
        OpRegistry::Global(), proto);
    OptimizerOptions opts;
    device_mgr_ = std::make_unique<StaticDeviceMgr>(std::move(devices));

      /*function_name=*/name, /*in_def=*/{}, /*out_def=*/{"out: float"},
      /*attr_def*/
      {}, /*node_def=*/{FunctionDefHelper::Const("one", 1.0f)},
      /*ret_def=*/{{"out", "out:output:0"}});
  *fdef_lib.add_function() = std::move(func);
  return fdef_lib;
}

TEST_F(BuildXlaOpsTest, ControlDepsPreserved) {
  const char* kXlaDeviceName = "/job:worker/replica:0/task:0/device:XLA_CPU:0";