#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"

        // broadcast, i.e. set as the input to each parallel operation. This
        // allows code like "tf.constant(1.)" or "tf.reduce_sum(..., axis=1)"
        // (where the value starts on the host), without allowing other implicit
        // copies/broadcasts. Other implicit copies may be supported eventually,
        // but need special handling for gradients (gradient of copy-on is not

#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/grappler/clusters/single_machine.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/lib/core/status_test_util.h"

  EXPECT_EQ(TF_GraphNumFunctions(host_graph_), 1);
  EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 0, s_), 0);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
  EXPECT_EQ(TF_GraphGetFunctions(host_graph_, funcs, 1, s_), 1);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
  AssertEqual(func0, funcs[0]);
  TF_DeleteFunction(funcs[0]);

}

void ReplaceTaskInServerDef(tensorflow::ServerDef* server_def, int task_index,
                            const string& host, int port) {
  tensorflow::JobDef* job_def = server_def->mutable_cluster()->mutable_job(0);
  job_def->mutable_tasks()->at(task_index) =
      tensorflow::strings::StrCat(host, ":", port);
}

std::vector<std::string> ListDeviceNames(TFE_Context* ctx) {
  TF_Status* status = TF_NewStatus();

  (*sess_options.config.mutable_device_count())["CPU"] = num_cpus;
  status->status =
      tensorflow::DeviceFactory::AddCpuDevices(sess_options, prefix, &devices);

  // Remove the device that has the host device name since host device is alreay
  // in an initialized context.
  for (auto d = devices.begin(); d != devices.end();) {
    if (absl::StrContains(d->get()->name(), "CPU:0")) {
      d = devices.erase(d);
    } else {

    const tensorflow::ServerDef& cluster_server_def, int task_index) {
  tensorflow::ServerDef single_host_server_def;
  single_host_server_def.set_job_name("worker");
  single_host_server_def.set_protocol(cluster_server_def.protocol());
  single_host_server_def.set_task_index(0);
  tensorflow::ClusterDef* cluster_def =
      single_host_server_def.mutable_cluster();
  tensorflow::JobDef* job_def = cluster_def->add_job();

    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    TFE_TensorHandle* on_host =
        TFE_TensorHandleCopyToDevice(copy_aliased, ctx, "CPU:0", status);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    TF_Tensor* resolved = TFE_TensorHandleResolve(on_host, status);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    const float* resolved_data =