chore: import upstream snapshot with attribution

tensorflow/c/README.md

@@ -0,0 +1,7 @@
+# TensorFlow C API
+
+- See [www.tensorflow.org/install/lang_c](https://www.tensorflow.org/install/lang_c)
+- Nightly builds:
+  - [Linux CPU-only](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-cpu-linux-x86_64.tar.gz)
+  - [Linux GPU](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-gpu-linux-x86_64.tar.gz)
+  - [MacOS CPU-only](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-cpu-darwin-x86_64.tar.gz)

tensorflow/c/c_api_experimental.cc

@@ -0,0 +1,777 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/c_api_experimental.h"
+
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/substitute.h"
+#include "absl/synchronization/notification.h"
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/checkpoint_reader.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/eager/tfe_op_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_buffer_internal.h"
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/core/common_runtime/eager/attr_builder.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/common_runtime/pluggable_device/pluggable_device_plugin_init.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/framework/collective.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/platform/blocking_counter.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/net.h"
+#include "tensorflow/core/platform/platform.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+using tensorflow::FunctionDef;
+using tensorflow::Node;
+using tensorflow::NodeBuilder;
+using tensorflow::Status;
+
+namespace {
+typedef std::unique_ptr<TF_Function, decltype(&TF_DeleteFunction)>
+    UniqueFuncPtr;
+}
+
+// struct TF_Operation { tensorflow::Node node; };
+static TF_Operation* ToTF_Operation(Node* node) {
+  return static_cast<TF_Operation*>(static_cast<void*>(node));
+}
+
+void TF_EnableXLACompilation(TF_SessionOptions* options, unsigned char enable) {
+  tensorflow::ConfigProto& config = options->options.config;
+  auto* optimizer_options =
+      config.mutable_graph_options()->mutable_optimizer_options();
+  if (enable) {
+    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::ON_1);
+
+    // These XLA flags are needed to trigger XLA properly from C (more generally
+    // non-Python) clients. If this API is called again with `enable` set to
+    // false, it is safe to keep these flag values as is.
+    tensorflow::MarkForCompilationPassFlags* flags =
+        tensorflow::GetMarkForCompilationPassFlags();
+    flags->tf_xla_cpu_global_jit = true;
+    flags->tf_xla_min_cluster_size = 1;
+  } else {
+    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF);
+  }
+}
+
+unsigned char TF_SetXlaEnableLazyCompilation(unsigned char enable) {
+  tensorflow::BuildXlaOpsPassFlags* flags =
+      tensorflow::GetBuildXlaOpsPassFlags();
+  bool original = flags->tf_xla_enable_lazy_compilation;
+  flags->tf_xla_enable_lazy_compilation = enable;
+  return original;
+}
+
+unsigned char TF_SetTfXlaCpuGlobalJit(unsigned char enable) {
+  tensorflow::MarkForCompilationPassFlags* flags =
+      tensorflow::GetMarkForCompilationPassFlags();
+  bool original = flags->tf_xla_cpu_global_jit;
+  flags->tf_xla_cpu_global_jit = static_cast<bool>(enable);
+  return static_cast<unsigned char>(original);
+}
+
+void TF_SetXlaAutoJitMode(const char* mode) {
+  tensorflow::SetXlaAutoJitFlagFromFlagString(mode);
+}
+
+unsigned char TF_GetXlaAutoJitEnabled() {
+  tensorflow::XlaAutoJitFlag flag =
+      tensorflow::GetMarkForCompilationPassFlags()->xla_auto_jit_flag;
+  return static_cast<unsigned char>(flag.optimization_level_single_gpu > 0 ||
+                                    flag.optimization_level_general > 0);
+}
+
+unsigned char TF_GetXlaConstantFoldingDisabled() {
+  return static_cast<unsigned char>(
+      tensorflow::GetBuildXlaOpsPassFlags()->tf_xla_disable_constant_folding);
+}
+
+void TF_SetXlaConstantFoldingDisabled(unsigned char should_enable) {
+  tensorflow::GetBuildXlaOpsPassFlags()->tf_xla_disable_constant_folding =
+      static_cast<bool>(should_enable);
+}
+
+void TF_SetXlaMinClusterSize(int size) {
+  tensorflow::MarkForCompilationPassFlags* flags =
+      tensorflow::GetMarkForCompilationPassFlags();
+  flags->tf_xla_min_cluster_size = size;
+}
+
+TF_Buffer* TF_CreateConfig(unsigned char enable_xla_compilation,
+                           unsigned char gpu_memory_allow_growth,
+                           unsigned int num_cpu_devices) {
+  tensorflow::ConfigProto config;
+  auto* optimizer_options =
+      config.mutable_graph_options()->mutable_optimizer_options();
+  if (enable_xla_compilation) {
+    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::ON_1);
+
+    // These XLA flags are needed to trigger XLA properly from C (more generally
+    // non-Python) clients. If this API is called again with `enable` set to
+    // false, it is safe to keep these flag values as is.
+    tensorflow::MarkForCompilationPassFlags* flags =
+        tensorflow::GetMarkForCompilationPassFlags();
+    flags->tf_xla_cpu_global_jit = true;
+    flags->tf_xla_min_cluster_size = 1;
+  } else {
+    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF);
+  }
+
+  auto* gpu_options = config.mutable_gpu_options();
+  gpu_options->set_allow_growth(gpu_memory_allow_growth);
+
+  (*config.mutable_device_count())["CPU"] = num_cpu_devices;
+
+  // TODO(b/113217601): This is needed for EagerContext::runner_ to use a
+  // threadpool, so that we avoid the possibility of running the runner_ in the
+  // threadpool of GPU event mgr, as that can trigger more callbacks to be
+  // scheduled on that same threadpool, causing a deadlock in cases where the
+  // caller of event_mgr->ThenExecute() blocks on the completion of the callback
+  // (as in the case of ConstOp kernel creation on GPU, which involves copying a
+  // CPU tensor to GPU).
+  // Setting a larger thread pool does not help with the Swift caller, as we use
+  // a different TFE context for each thread of execution (for running graph
+  // functions, and their send/recvs corountines).
+  config.set_inter_op_parallelism_threads(1);
+
+  TF_Buffer* ret = TF_NewBuffer();
+  TF_CHECK_OK(MessageToBuffer(config, ret));
+  return ret;
+}
+
+TF_Buffer* TF_CreateRunOptions(unsigned char enable_full_trace) {
+  tensorflow::RunOptions options;
+  if (enable_full_trace) {
+    options.set_trace_level(tensorflow::RunOptions::FULL_TRACE);
+  } else {
+    options.set_trace_level(tensorflow::RunOptions::NO_TRACE);
+  }
+  TF_Buffer* ret = TF_NewBuffer();
+  TF_CHECK_OK(MessageToBuffer(options, ret));
+  return ret;
+}
+
+const char* TF_GraphDebugString(TF_Graph* graph, size_t* len) {
+  tensorflow::mutex_lock c(graph->mu);
+  const auto& debug_str = graph->graph.ToGraphDefDebug().DebugString();
+  *len = debug_str.size();
+  char* ret = static_cast<char*>(malloc(*len + 1));
+  memcpy(ret, debug_str.c_str(), *len + 1);
+  return ret;
+}
+
+char* TF_FunctionDebugString(TF_Function* func, size_t* len) {
+  const auto& debug_str = DebugString(func->record->fdef());
+  *len = debug_str.size();
+  char* ret = static_cast<char*>(malloc(*len + 1));
+  memcpy(ret, debug_str.c_str(), *len + 1);
+  return ret;
+}
+
+// On success, returns a set of TF_Function instances from `text_proto` of
+// GraphDef type. These functions must be deleted by calling TF_DeleteFunction.
+//
+// If `mutate_proto_func` is non-NULL, run it over each FunctionDef proto,
+// before creating a TF_Function out of the possibly mutated proto.
+static std::vector<UniqueFuncPtr> CreateFunctionsFromTextProto(
+    const char* text_proto,
+    std::function<void(FunctionDef*)>* mutate_proto_func, TF_Status* status) {
+  tensorflow::GraphDef gdef;
+  if (!tensorflow::protobuf::TextFormat::ParseFromString(text_proto, &gdef)) {
+    status->status = absl::InternalError(
+        absl::StrCat("Invalid text proto for GraphDef: ", text_proto));
+    return {};
+  }
+  const auto& fdef_lib = gdef.library();
+  if (fdef_lib.gradient_size() > 0) {
+    status->status = absl::InternalError(absl::StrCat(
+        "GradientDef is not supported in reading Dataset related functions: ",
+        text_proto));
+    return {};
+  }
+  std::vector<UniqueFuncPtr> ret;
+  for (const FunctionDef& fdef : fdef_lib.function()) {
+    // Make a copy so that we can mutate it.
+    FunctionDef fdef_to_load = fdef;
+    if (mutate_proto_func) {
+      (*mutate_proto_func)(&fdef_to_load);
+    }
+    VLOG(1) << "Adding func to graph: " << fdef_to_load.DebugString();
+    std::vector<char> binary_proto_buf(fdef_to_load.ByteSizeLong());
+    fdef_to_load.SerializeToArray(binary_proto_buf.data(),
+                                  binary_proto_buf.size());
+    TF_Function* func = TF_FunctionImportFunctionDef(
+        binary_proto_buf.data(), binary_proto_buf.size(), status);
+    if (!status->status.ok()) return {};
+    ret.push_back(UniqueFuncPtr(func, TF_DeleteFunction));
+  }
+  return ret;
+}
+
+TF_Tensor* TF_DequeueNamedTensor(TF_Session* session, int tensor_id,
+                                 TF_Status* status) {
+  assert(session);
+  {
+    tensorflow::mutex_lock c(session->graph->mu);
+    VLOG(1) << "Dequeuing named tensor with id " << tensor_id
+            << ", with input graph: "
+            << session->graph->graph.ToGraphDefDebug().DebugString();
+  }
+
+  TF_Operation* dequeue_op = TF_GraphOperationByName(
+      session->graph, absl::StrCat("fifo_queue_dequeue_", tensor_id).c_str());
+  if (dequeue_op == nullptr) {
+    status->status =
+        absl::InternalError("Unable to find the dequeue node in the TF graph.");
+    return nullptr;
+  }
+
+  VLOG(1) << "Running the dequeue op";
+  TF_Output output{dequeue_op, 0};
+  TF_Tensor* ret;
+  TF_SessionRun(session, /*run_options*/ nullptr,
+                // input related parameters
+                /*inputs*/ nullptr, /*input_values*/ nullptr, /*ninputs*/ 0,
+                // output related parameters
+                /*outputs*/ &output, /*output_values*/ &ret,
+                /*noutputs*/ 1,
+                /*targets*/ nullptr, /*ntargets*/ 0,
+                /*run_metadata*/ nullptr, status);
+  if (VLOG_IS_ON(1) && status->status.ok()) {
+    tensorflow::Tensor tensor;
+    if (tensorflow::TF_TensorToTensor(ret, &tensor).ok()) {
+      VLOG(1) << "Dequeued tensor content: " << tensor.DebugString();
+    }
+  }
+  return ret;
+}
+
+void TF_EnqueueNamedTensor(TF_Session* session, int tensor_id,
+                           TF_Tensor* tensor, TF_Status* status) {
+  assert(session);
+  {
+    tensorflow::mutex_lock c(session->graph->mu);
+    if (VLOG_IS_ON(1)) {
+      VLOG(1) << "Enqueuing named tensor with id " << tensor_id
+              << ", with input graph: "
+              << session->graph->graph.ToGraphDefDebug().DebugString();
+      tensorflow::Tensor internal_tensor;
+      if (tensorflow::TF_TensorToTensor(tensor, &internal_tensor).ok()) {
+        VLOG(1) << "Enqueu'ing tensor content: "
+                << internal_tensor.DebugString();
+      }
+    }
+  }
+
+  TF_Operation* enqueue_op = TF_GraphOperationByName(
+      session->graph, absl::StrCat("fifo_queue_enqueue_", tensor_id).c_str());
+  if (enqueue_op == nullptr) {
+    status->status =
+        absl::InternalError("Unable to find the enqueue node in the TF graph.");
+    return;
+  }
+
+  TF_Operation* placeholder_op = TF_GraphOperationByName(
+      session->graph, absl::StrCat("arg_tensor_enqueue_", tensor_id).c_str());
+  if (placeholder_op == nullptr) {
+    status->status = absl::InternalError(
+        "Unable to find the placeholder node as input to enqueue in the TF "
+        "graph.");
+    return;
+  }
+
+  VLOG(1) << "Running the enqueue op";
+  TF_Output input{placeholder_op, 0};
+  TF_SessionRun(session, /*run_options*/ nullptr,
+                // input related parameters
+                /*inputs*/ &input, /*input_values*/ &tensor, /*ninputs*/ 1,
+                // output related parameters
+                /*outputs*/ nullptr, /*output_values*/ nullptr, /*noutputs*/ 0,
+                /*targets*/ &enqueue_op, /*ntargets*/ 1,
+                /*run_metadata*/ nullptr, status);
+  VLOG(1) << "Enqueuing is done.";
+}
+
+TF_Buffer* TFE_GetServerDef(const char* text_proto, TF_Status* status) {
+  tensorflow::ServerDef server_def;
+  if (!tensorflow::protobuf::TextFormat::ParseFromString(text_proto,
+                                                         &server_def)) {
+    status->status = absl::InternalError(
+        absl::StrCat("Invalid text proto for ServerDef: ", text_proto));
+    return nullptr;
+  }
+  status->status = absl::Status();
+  TF_Buffer* ret = TF_NewBuffer();
+  TF_CHECK_OK(MessageToBuffer(server_def, ret));
+  return ret;
+}
+
+void TF_MakeInternalErrorStatus(TF_Status* status, const char* errMsg) {
+  status->status = absl::InternalError(errMsg);
+}
+
+struct TF_CheckpointReader : public tensorflow::checkpoint::CheckpointReader {
+  using tensorflow::checkpoint::CheckpointReader::CheckpointReader;
+  std::vector<std::string> variable_list;
+};
+
+TF_CheckpointReader* TF_NewCheckpointReader(const char* filename,
+                                            TF_Status* status) {
+  TF_CheckpointReader* reader = new TF_CheckpointReader(filename, status);
+  if (!status->status.ok()) {
+    TF_DeleteCheckpointReader(reader);
+    return nullptr;
+  }
+  const auto& m = reader->GetVariableToDataTypeMap();
+  for (auto it = m.begin(); it != m.end(); ++it)
+    reader->variable_list.push_back(it->first);
+  std::sort(reader->variable_list.begin(), reader->variable_list.end());
+  return reader;
+}
+
+void TF_DeleteCheckpointReader(TF_CheckpointReader* reader) { delete reader; }
+
+int TF_CheckpointReaderHasTensor(TF_CheckpointReader* reader,
+                                 const char* name) {
+  return reader->HasTensor(name);
+}
+
+const char* TF_CheckpointReaderGetVariable(TF_CheckpointReader* reader,
+                                           int index) {
+  return reader->variable_list[index].c_str();
+}
+
+int TF_CheckpointReaderSize(TF_CheckpointReader* reader) {
+  return reader->variable_list.size();
+}
+
+TF_DataType TF_CheckpointReaderGetVariableDataType(TF_CheckpointReader* reader,
+                                                   const char* name) {
+  const auto& m = reader->GetVariableToDataTypeMap();
+  return static_cast<TF_DataType>(m.at(name));
+}
+
+TF_Tensor* TF_CheckpointReaderGetTensor(TF_CheckpointReader* reader,
+                                        const char* name, TF_Status* status) {
+  std::unique_ptr<tensorflow::Tensor> tensor;
+  reader->GetTensor(name, &tensor, status);
+  if (!status->status.ok()) return nullptr;
+  return tensorflow::TF_TensorFromTensor(*tensor, &status->status);
+}
+
+void TF_CheckpointReaderGetVariableShape(TF_CheckpointReader* reader,
+                                         const char* name, int64_t* dims,
+                                         int num_dims, TF_Status* status) {
+  const auto& shape = reader->GetVariableToShapeMap().at(name);
+  int rank = shape.dims();
+  if (num_dims != rank) {
+    status->status = absl::InvalidArgumentError(absl::StrCat(
+        "Expected rank is ", num_dims, " but actual rank is ", rank));
+    return;
+  }
+  for (int i = 0; i < num_dims; i++) {
+    dims[i] = shape.dim_size(i);
+  }
+}
+
+int TF_CheckpointReaderGetVariableNumDims(TF_CheckpointReader* reader,
+                                          const char* name) {
+  const auto& m = reader->GetVariableToShapeMap();
+  return m.at(name).dims();
+}
+
+// This builder is used in the eager API to build a NodeDef.
+struct TF_AttrBuilder : public tensorflow::AttrBuilder {
+  using tensorflow::AttrBuilder::AttrBuilder;
+  // The string buffers to make sure that any `attr_name` we pass into
+  // `builder->Set()` will outlive the subsequent
+  // `TF_AttrBuilderCheckCanRunOnDevice()` call(s) on the same `builder`.
+  std::set<std::string> attr_names;
+};
+
+TF_AttrBuilder* TF_NewAttrBuilder(const char* op_name) {
+  return new TF_AttrBuilder(op_name);
+}
+
+void TF_DeleteAttrBuilder(TF_AttrBuilder* builder) { delete builder; }
+
+void TF_AttrBuilderSetType(TF_AttrBuilder* builder, const char* attr_name,
+                           TF_DataType value) {
+  auto iter = builder->attr_names.insert(attr_name).first;
+  builder->Set(*iter, static_cast<tensorflow::DataType>(value));
+}
+
+void TF_AttrBuilderSetTypeList(TF_AttrBuilder* builder, const char* attr_name,
+                               const TF_DataType* values, int num_values) {
+  auto iter = builder->attr_names.insert(attr_name).first;
+  builder->Set(*iter, tensorflow::gtl::ArraySlice<const tensorflow::DataType>(
+                          reinterpret_cast<const tensorflow::DataType*>(values),
+                          num_values));
+}
+
+void TF_AttrBuilderCheckCanRunOnDevice(TF_AttrBuilder* builder,
+                                       const char* device_type,
+                                       TF_Status* status) {
+  status->status = tensorflow::FindKernelDef(
+      tensorflow::DeviceType(device_type), builder->BuildNodeDef(),
+      /* def = */ nullptr, /* kernel_class_name = */ nullptr);
+}
+
+const char* TF_GetNumberAttrForOpListInput(const char* op_name, int input_index,
+                                           TF_Status* status) {
+  const tensorflow::OpDef* op_def = nullptr;
+  status->status =
+      tensorflow::OpRegistry::Global()->LookUpOpDef(op_name, &op_def);
+  if (!status->status.ok()) return nullptr;
+
+  if (input_index >= op_def->input_arg_size() || input_index < 0) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat(input_index, " out of range for ", op_name));
+    return nullptr;
+  }
+
+  const tensorflow::OpDef_ArgDef& input_arg = op_def->input_arg()[input_index];
+
+  if (input_arg.number_attr().empty()) {
+    status->status = absl::NotFoundError(
+        absl::StrCat(op_name, " does not have number_attr() defined."));
+    return nullptr;
+  }
+
+  // The returned string is owned by OpRegistry, so liveness is not a concern.
+  return input_arg.number_attr().c_str();
+}
+
+int TF_OpIsStateful(const char* op_type, TF_Status* status) {
+  const tensorflow::OpRegistrationData* op_reg_data;
+  status->status =
+      tensorflow::OpRegistry::Global()->LookUp(op_type, &op_reg_data);
+  if (!status->status.ok()) {
+    return 0;
+  }
+  return op_reg_data->op_def.is_stateful();
+}
+
+void TF_InitMain(const char* usage, int* argc, char*** argv) {
+  tensorflow::port::InitMain(usage, argc, argv);
+}
+
+int TF_PickUnusedPortOrDie() { return tsl::net::PickUnusedPortOrDie(); }
+
+TFE_TensorHandle* TFE_NewTensorHandleFromScalar(TF_DataType data_type,
+                                                void* data, size_t len,
+                                                TF_Status* status) {
+  auto dtype = static_cast<tensorflow::DataType>(data_type);
+  if (!tensorflow::DataTypeCanUseMemcpy(dtype)) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("DataType ", dtype, " cannot use memcpy"));
+    return nullptr;
+  }
+
+  size_t expected_len = static_cast<size_t>(tensorflow::DataTypeSize(dtype));
+  if (len > expected_len) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("len (", len, ") is larger than expected scalar size for ",
+                     dtype, " (", expected_len, " bytes)"));
+    return nullptr;
+  }
+
+  tensorflow::Tensor tensor(dtype, tensorflow::TensorShape({}));
+  std::memcpy(tensorflow::TensorCApi::Buffer(tensor)->data(), data, len);
+
+  status->status = absl::OkStatus();
+  return tensorflow::wrap(tensorflow::TensorHandle::CreateLocalHandle(tensor));
+}
+
+// Set server_def on the context, possibly updating it.
+TF_CAPI_EXPORT extern void TFE_EnableCollectiveOps(TFE_Context* ctx,
+                                                   const void* proto,
+                                                   size_t proto_len,
+                                                   TF_Status* status) {
+  tensorflow::ServerDef server_def;
+  if (!server_def.ParseFromArray(proto, proto_len)) {
+    status->status = absl::InvalidArgumentError(
+        "Invalid tensorflow.ServerDef protocol buffer");
+    return;
+  }
+  status->status = tensorflow::unwrap(ctx)->EnableCollectiveOps(server_def);
+}
+
+TF_CAPI_EXPORT extern void TFE_AbortCollectiveOps(TFE_Context* ctx,
+                                                  TF_Status* status) {
+  tensorflow::EagerContext* context =
+      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
+  auto collective_executor_handle = context->GetCollectiveExecutorHandle();
+  collective_executor_handle->get()->StartAbort(status->status);
+}
+
+TF_CAPI_EXPORT extern void TFE_CollectiveOpsCheckPeerHealth(
+    TFE_Context* ctx, const char* task, int64_t timeout_in_ms,
+    TF_Status* status) {
+  tensorflow::EagerContext* context =
+      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
+  auto collective_executor_handle = context->GetCollectiveExecutorHandle();
+  absl::Notification done;
+  collective_executor_handle->get()->remote_access()->CheckPeerHealth(
+      task, timeout_in_ms, [&done, status](const Status& s) {
+        status->status = s;
+        done.Notify();
+      });
+  done.WaitForNotification();
+}
+
+TF_ShapeAndTypeList* TF_NewShapeAndTypeList(int num_items) {
+  TF_ShapeAndTypeList* result = new TF_ShapeAndTypeList;
+  result->num_items = num_items;
+  result->items = (num_items == 0) ? nullptr : new TF_ShapeAndType[num_items]();
+  return result;
+}
+
+void TF_ShapeAndTypeListSetShape(TF_ShapeAndTypeList* shape_list, int index,
+                                 const int64_t* dims, int num_dims) {
+  CHECK(index >= 0 && index < shape_list->num_items);
+  TF_ShapeAndType& shape = shape_list->items[index];
+  CHECK(shape.dims == nullptr) << "Shape at " << index << " is already set!";
+  CHECK(num_dims >= 0) << "Number of dimensions cannot be negative!";
+  shape.num_dims = num_dims;
+  shape.dims = new int64_t[num_dims];
+  memcpy(shape.dims, dims, sizeof(int64_t) * num_dims);
+}
+
+void TF_ShapeAndTypeListSetUnknownShape(TF_ShapeAndTypeList* shape_list,
+                                        int index) {
+  CHECK(index >= 0 && index < shape_list->num_items);
+  TF_ShapeAndType& shape = shape_list->items[index];
+  CHECK(shape.dims == nullptr) << "Shape at " << index << " is already set!";
+  shape.num_dims = -1;
+  shape.dims = nullptr;
+}
+
+void TF_ShapeAndTypeListSetDtype(TF_ShapeAndTypeList* shape_list, int index,
+                                 TF_DataType dtype) {
+  CHECK(index >= 0 && index < shape_list->num_items);
+  TF_ShapeAndType& shape_and_type = shape_list->items[index];
+  shape_and_type.dtype = dtype;
+}
+
+void TF_DeleteShapeAndTypeList(TF_ShapeAndTypeList* shape_list) {
+  if (shape_list == nullptr) return;
+  for (size_t i = 0; i < shape_list->num_items; ++i) {
+    delete[] shape_list->items[i].dims;
+  }
+  delete[] shape_list->items;
+  delete shape_list;
+}
+
+void TF_DeleteShapeAndTypeListArray(TF_ShapeAndTypeList** shape_list_array,
+                                    int num_items) {
+  if (shape_list_array == nullptr) return;
+  for (int i = 0; i < num_items; ++i) {
+    TF_DeleteShapeAndTypeList(shape_list_array[i]);
+  }
+  delete[] shape_list_array;
+}
+
+namespace tensorflow {
+absl::Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
+
+// Helpers for loadding a TensorFlow PluggableDevice plugin (a .so file).
+absl::Status LoadPluggableDeviceLibrary(const char* library_filename,
+                                        void** result);
+}  // namespace tensorflow
+
+void TFE_InferShapes(TFE_Op* tfe_op, TF_ShapeAndTypeList* input_shapes,
+                     TF_Tensor** input_tensors,
+                     TF_ShapeAndTypeList* input_tensors_as_shapes,
+                     TF_ShapeAndTypeList** input_resource_shapes_and_types,
+                     TF_ShapeAndTypeList** output_shapes,
+                     TF_ShapeAndTypeList*** output_resource_shapes_and_types,
+                     TF_Status* status) {
+  using tensorflow::NodeDef;
+  using tensorflow::OpRegistrationData;
+  using tensorflow::Tensor;
+  using tensorflow::shape_inference::DimensionHandle;
+  using tensorflow::shape_inference::InferenceContext;
+  using tensorflow::shape_inference::ShapeAndType;
+  using tensorflow::shape_inference::ShapeHandle;
+
+  const int num_inputs = input_shapes->num_items;
+  NodeDef node_def;
+  tensorflow::ImmediateExecutionOperation* op = tensorflow::unwrap(tfe_op);
+  node_def.set_name(op->Name());
+  node_def.set_op(op->Name());
+  for (int i = 0; i < num_inputs; ++i) {
+    node_def.add_input("dummy_input");
+  }
+  OperationFromInterface(op)->Attrs().FillAttrValueMap(node_def.mutable_attr());
+
+  const tensorflow::OpRegistrationData* op_reg_data;
+  status->status =
+      tensorflow::OpRegistry::Global()->LookUp(node_def.op(), &op_reg_data);
+  if (!status->status.ok()) return;
+
+  // Initialize a input_tensor vector with `nullptr` values.
+  std::vector<const Tensor*> input_tensors_vector(num_inputs, nullptr);
+  // A vector to keep track of newly created `tf::Tensor` objects.
+  std::vector<Tensor> all_input_tensors;
+  // Update the vector with information from `input_tensors` if provided.
+  if (input_tensors != nullptr) {
+    // Note that we take the address of the elements in `all_input_tensors`
+    // below. Allocate enough space so that no reallocation happens, which will
+    // make the pointers invalid.
+    all_input_tensors.reserve(num_inputs);
+    for (int i = 0; i < num_inputs; ++i) {
+      if (input_tensors[i] == nullptr) continue;
+      all_input_tensors.emplace_back();
+      Tensor& input_tensor = all_input_tensors.back();
+      status->status = TF_TensorToTensor(input_tensors[i], &input_tensor);
+      if (!status->status.ok()) return;
+      input_tensors_vector[i] = &input_tensor;
+    }
+  }
+
+  // Create an inference context with dummy values, which will be updated later.
+  InferenceContext c(TF_GRAPH_DEF_VERSION, node_def, op_reg_data->op_def,
+                     std::vector<ShapeHandle>(num_inputs), input_tensors_vector,
+                     {},
+                     std::vector<std::unique_ptr<std::vector<ShapeAndType>>>());
+
+  // Set input_shapes.
+  for (int i = 0; i < num_inputs; ++i) {
+    std::vector<DimensionHandle> dims;
+    const TF_ShapeAndType& input_shape = input_shapes->items[i];
+    if (input_shape.num_dims == InferenceContext::kUnknownRank) {
+      c.SetInput(i, c.UnknownShape());
+      continue;
+    }
+    dims.reserve(input_shape.num_dims);
+    for (int j = 0; j < input_shape.num_dims; ++j) {
+      dims.push_back(c.MakeDim(input_shape.dims[j]));
+    }
+    c.SetInput(i, c.MakeShape(dims));
+  }
+
+  // TODO(bgogul): Handle input_tensors_as_shapes.
+  // TODO(bgogul): Handle input_resource_shapes_and_types.
+
+  status->status = c.construction_status();
+  if (!status->status.ok()) return;
+
+  if (op_reg_data->shape_inference_fn == nullptr) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("No shape inference function exists for op '",
+                     node_def.op(), "', did you forget to define it?"));
+    return;
+  }
+
+  status->status = c.Run(op_reg_data->shape_inference_fn);
+  if (!status->status.ok()) return;
+
+  // Set output_shapes.
+  TF_ShapeAndTypeList* output_shapes_result =
+      TF_NewShapeAndTypeList(c.num_outputs());
+  for (int i = 0; i < c.num_outputs(); ++i) {
+    ShapeHandle shape_handle = c.output(i);
+    TF_ShapeAndType& shape = output_shapes_result->items[i];
+    shape.num_dims = c.Rank(shape_handle);
+    if (shape.num_dims == InferenceContext::kUnknownRank) {
+      shape.dims = nullptr;
+      continue;
+    }
+    shape.dims = new int64_t[shape.num_dims];
+    for (size_t j = 0; j < shape.num_dims; ++j) {
+      shape.dims[j] = c.Value(c.Dim(shape_handle, j));
+    }
+  }
+  if (output_shapes != nullptr) *output_shapes = output_shapes_result;
+
+  // TODO(bgogul): Set output_resource_shapes_and_types.
+}
+
+void TF_ImportGraphDefOptionsSetValidateColocationConstraints(
+    TF_ImportGraphDefOptions* opts, unsigned char enable) {
+  opts->opts.validate_colocation_constraints = enable;
+}
+
+// Load a Pluggable Device library.
+// On success, returns the handle to library in result and return OK from the
+// function. Otherwise return nullptr in result and error Status from the
+// function.
+//
+// If `library_filename` has already been loaded, we return a cached handle.
+// Device and Kernels/Ops are registered as globals when a library is loaded
+// for the first time.
+TF_Library* TF_LoadPluggableDeviceLibrary(const char* library_filename,
+                                          TF_Status* status) {
+#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD)
+  status->status = absl::UnimplementedError(
+      "PluggableDevice plugin functionality is not supported on mobile");
+  return nullptr;
+#else
+  TF_Library* lib_handle = new TF_Library;
+  static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED);
+  static std::unordered_map<std::string, void*>* loaded_libs =
+      new std::unordered_map<std::string, void*>();
+  tensorflow::Env* env = tensorflow::Env::Default();
+  {
+    tensorflow::mutex_lock lock(mu);
+    auto it = loaded_libs->find(library_filename);
+    if (it != loaded_libs->end()) {
+      lib_handle->lib_handle = it->second;
+    } else {
+      status->status =
+          env->LoadDynamicLibrary(library_filename, &lib_handle->lib_handle);
+      if (status->status.ok()) {
+        TF_CHECK_OK(
+            tensorflow::RegisterPluggableDevicePlugin(lib_handle->lib_handle));
+      } else {
+        delete lib_handle;
+        return nullptr;
+      }
+    }
+    return lib_handle;
+  }
+#endif
+}
+
+void TF_DeletePluggableDeviceLibraryHandle(TF_Library* lib_handle) {
+  delete lib_handle;
+}
+
+void TF_GraphRemoveFunction(TF_Graph* g, const char* func_name,
+                            TF_Status* status) {
+  tensorflow::mutex_lock l(g->mu);
+  status->status = g->graph.mutable_flib_def()->RemoveFunction(func_name);
+}

tensorflow/c/c_api_experimental.h

@@ -0,0 +1,324 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_API_EXPERIMENTAL_H_
+#define TENSORFLOW_C_C_API_EXPERIMENTAL_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/eager/c_api.h"
+
+// --------------------------------------------------------------------------
+// Experimental C API for TensorFlow.
+//
+// The API here is subject to changes in the future.
+// --------------------------------------------------------------------------
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// When `enable` is true, set
+// tensorflow.ConfigProto.OptimizerOptions.global_jit_level to ON_1, and also
+// set XLA flag values to prepare for XLA compilation. Otherwise set
+// global_jit_level to OFF.
+//
+// This and the next API are syntax sugar over TF_SetConfig(), and is used by
+// clients that cannot read/write the tensorflow.ConfigProto proto.
+// TODO: Migrate to TF_CreateConfig() below.
+TF_CAPI_EXPORT extern void TF_EnableXLACompilation(TF_SessionOptions* options,
+                                                   unsigned char enable);
+
+// Set XLA's internal BuildXlaOpsPassFlags.tf_xla_enable_lazy_compilation to the
+// value of 'enabled'. Also returns the original value of that flag.
+//
+// Use in tests to allow XLA to fallback to TF classic. This has global effect.
+TF_CAPI_EXPORT unsigned char TF_SetXlaEnableLazyCompilation(
+    unsigned char enable);
+TF_CAPI_EXPORT unsigned char TF_SetTfXlaCpuGlobalJit(unsigned char enable);
+
+// Sets XLA's auto jit mode according to the specified string, which is parsed
+// as if passed in XLA_FLAGS. This has global effect.
+TF_CAPI_EXPORT void TF_SetXlaAutoJitMode(const char* mode);
+
+// Returns whether the single GPU or general XLA auto jit optimizations are
+// enabled through MarkForCompilationPassFlags.
+TF_CAPI_EXPORT unsigned char TF_GetXlaAutoJitEnabled();
+
+// Sets XLA's minimum cluster size. This has global effect.
+TF_CAPI_EXPORT void TF_SetXlaMinClusterSize(int size);
+
+// Gets/Sets TF/XLA flag for whether(true) or not(false) to disable constant
+// folding. This is for testing to ensure that XLA is being tested rather than
+// Tensorflow's CPU implementation through constant folding.
+TF_CAPI_EXPORT unsigned char TF_GetXlaConstantFoldingDisabled();
+TF_CAPI_EXPORT void TF_SetXlaConstantFoldingDisabled(
+    unsigned char should_enable);
+
+// Create a serialized tensorflow.ConfigProto proto, where:
+//
+// a) ConfigProto.optimizer_options.global_jit_level is set to ON_1 if
+// `enable_xla_compilation` is non-zero, and OFF otherwise.
+// b) ConfigProto.gpu_options.allow_growth is set to `gpu_memory_allow_growth`.
+// c) ConfigProto.device_count is set to `num_cpu_devices`.
+TF_CAPI_EXPORT extern TF_Buffer* TF_CreateConfig(
+    unsigned char enable_xla_compilation, unsigned char gpu_memory_allow_growth,
+    unsigned int num_cpu_devices);
+
+// Create a serialized tensorflow.RunOptions proto, where RunOptions.trace_level
+// is set to FULL_TRACE if `enable_full_trace` is non-zero, and NO_TRACE
+// otherwise.
+TF_CAPI_EXPORT extern TF_Buffer* TF_CreateRunOptions(
+    unsigned char enable_full_trace);
+
+// Returns the graph content in a human-readable format, with length set in
+// `len`. The format is subject to change in the future.
+// The returned string is heap-allocated, and caller should call free() on it.
+TF_CAPI_EXPORT extern const char* TF_GraphDebugString(TF_Graph* graph,
+                                                      size_t* len);
+
+// Returns the function content in a human-readable format, with length set in
+// `len`. The format is subject to change in the future.
+// The returned string is heap-allocated, and caller should call free() on it.
+//
+// Do not return const char*, because some foreign language binding
+// (e.g. swift) cannot then call free() on the returned pointer.
+TF_CAPI_EXPORT extern char* TF_FunctionDebugString(TF_Function* func,
+                                                   size_t* len);
+
+// On success, dequeues a tensor from a TF-managed FifoQueue given by
+// `tensor_id`, associated with `session`. There must be a graph node named
+// "fifo_queue_dequeue_<tensor_id>", to be executed by this API call.
+
+// Caller must call TF_DeleteTensor() over the returned tensor. If the queue is
+// empty, this call is blocked.
+//
+// Tensors are enqueued via the corresponding TF enqueue op.
+// TODO(hongm): Add support for `timeout_ms`.
+TF_CAPI_EXPORT extern TF_Tensor* TF_DequeueNamedTensor(TF_Session* session,
+                                                       int tensor_id,
+                                                       TF_Status* status);
+
+// On success, enqueues `tensor` into a TF-managed FifoQueue given by
+// `tensor_id`, associated with `session`. There must be a graph node named
+// "fifo_queue_enqueue_<tensor_id>", to be executed by this API call. It reads
+// from a placeholder node "arg_tensor_enqueue_<tensor_id>".
+//
+// `tensor` is still owned by the caller. This call will be blocked if the queue
+// has reached its capacity, and will be unblocked when the queued tensors again
+// drop below the capacity due to dequeuing.
+//
+// Tensors are dequeued via the corresponding TF dequeue op.
+// TODO(hongm): Add support for `timeout_ms`.
+TF_CAPI_EXPORT extern void TF_EnqueueNamedTensor(TF_Session* session,
+                                                 int tensor_id,
+                                                 TF_Tensor* tensor,
+                                                 TF_Status* status);
+// Create a serialized tensorflow.ServerDef proto.
+TF_Buffer* TFE_GetServerDef(const char* text_proto, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TF_MakeInternalErrorStatus(TF_Status* status,
+                                                      const char* errMsg);
+
+// TF_NewCheckpointReader() return the CheckpointReader that can be use to
+// investigate or load the variable from the checkpoint file
+typedef struct TF_CheckpointReader TF_CheckpointReader;
+TF_CAPI_EXPORT extern TF_CheckpointReader* TF_NewCheckpointReader(
+    const char* filename, TF_Status* status);
+TF_CAPI_EXPORT extern void TF_DeleteCheckpointReader(
+    TF_CheckpointReader* reader);
+TF_CAPI_EXPORT extern int TF_CheckpointReaderHasTensor(
+    TF_CheckpointReader* reader, const char* name);
+// Get the variable name at the given index
+TF_CAPI_EXPORT extern const char* TF_CheckpointReaderGetVariable(
+    TF_CheckpointReader* reader, int index);
+// Get the number of variable in the checkpoint
+TF_CAPI_EXPORT extern int TF_CheckpointReaderSize(TF_CheckpointReader* reader);
+// Get the DataType of a variable
+TF_CAPI_EXPORT extern TF_DataType TF_CheckpointReaderGetVariableDataType(
+    TF_CheckpointReader* reader, const char* name);
+// Read the shape of a variable and write to `dims`
+TF_CAPI_EXPORT extern void TF_CheckpointReaderGetVariableShape(
+    TF_CheckpointReader* reader, const char* name, int64_t* dims, int num_dims,
+    TF_Status* status);
+// Get the number of dimension of a variable
+TF_CAPI_EXPORT extern int TF_CheckpointReaderGetVariableNumDims(
+    TF_CheckpointReader* reader, const char* name);
+// Load the weight of a variable
+TF_CAPI_EXPORT extern TF_Tensor* TF_CheckpointReaderGetTensor(
+    TF_CheckpointReader* reader, const char* name, TF_Status* status);
+
+// TF_NewAttrBuilder() returns an object that you can set attributes on as
+// though it were an op. This allows querying properties of that op for
+// type-checking purposes like if the op will run on a particular device type.
+typedef struct TF_AttrBuilder TF_AttrBuilder;
+TF_CAPI_EXPORT extern TF_AttrBuilder* TF_NewAttrBuilder(const char* op_name);
+TF_CAPI_EXPORT extern void TF_DeleteAttrBuilder(TF_AttrBuilder* builder);
+TF_CAPI_EXPORT extern void TF_AttrBuilderSetType(TF_AttrBuilder* builder,
+                                                 const char* attr_name,
+                                                 TF_DataType value);
+TF_CAPI_EXPORT extern void TF_AttrBuilderSetTypeList(TF_AttrBuilder* builder,
+                                                     const char* attr_name,
+                                                     const TF_DataType* values,
+                                                     int num_values);
+
+// Checks the tensorflow::NodeDef built via the methods above to see if it can
+// run on device_type.
+TF_CAPI_EXPORT extern void TF_AttrBuilderCheckCanRunOnDevice(
+    TF_AttrBuilder* builder, const char* device_type, TF_Status* status);
+
+// For argument number input_index, fetch the corresponding number_attr that
+// needs to be updated with the argument length of the input list.
+// Returns nullptr if there is any problem like op_name is not found, or the
+// argument does not support this attribute type.
+TF_CAPI_EXPORT extern const char* TF_GetNumberAttrForOpListInput(
+    const char* op_name, int input_index, TF_Status* status);
+
+// Returns 1 if the op is stateful, 0 otherwise. The return value is undefined
+// if the status is not ok.
+TF_CAPI_EXPORT extern int TF_OpIsStateful(const char* op_type,
+                                          TF_Status* status);
+
+// Platform specific initialization routine. Very few platforms actually require
+// this to be called.
+TF_CAPI_EXPORT void TF_InitMain(const char* usage, int* argc, char*** argv);
+
+// Platform-specific implementation to return an unused port. (This should used
+// in tests only.)
+TF_CAPI_EXPORT int TF_PickUnusedPortOrDie(void);
+
+// Fast path method that makes constructing a single scalar tensor require less
+// overhead and copies.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandleFromScalar(
+    TF_DataType data_type, void* data, size_t len, TF_Status* status);
+
+// Specify the server_def that enables collective ops.
+// This is different to the above function in that it doesn't create remote
+// contexts, and remotely executing ops is not possible. It just enables
+// communication for collective ops.
+TF_CAPI_EXPORT extern void TFE_EnableCollectiveOps(TFE_Context* ctx,
+                                                   const void* proto,
+                                                   size_t proto_len,
+                                                   TF_Status* status);
+
+// Aborts all ongoing collectives with the specified status. After abortion,
+// subsequent collectives will error with this status immediately. To reset the
+// collectives, create a new EagerContext.
+//
+// This is intended to be used when a peer failure is detected.
+TF_CAPI_EXPORT extern void TFE_AbortCollectiveOps(TFE_Context* ctx,
+                                                  TF_Status* status);
+
+// Checks the health of collective ops peers. Explicit health check is needed in
+// multi worker collective ops to detect failures in the cluster.  If a peer is
+// down, collective ops may hang.
+TF_CAPI_EXPORT extern void TFE_CollectiveOpsCheckPeerHealth(
+    TFE_Context* ctx, const char* task, int64_t timeout_in_ms,
+    TF_Status* status);
+
+// Information about the shape of a Tensor and its type.
+struct TF_ShapeAndType {
+  // Number of dimensions. -1 indicates unknown rank.
+  int num_dims;
+  // Array of dimensions. -1 indicates unknown dim.
+  int64_t* dims;
+  // The data type. May be 0 to denote unknown type.
+  TF_DataType dtype;
+};
+
+typedef struct TF_ShapeAndType TF_ShapeAndType;
+
+// A list of TF_ShapeAndType elements..
+struct TF_ShapeAndTypeList {
+  int num_items;
+  TF_ShapeAndType* items;
+};
+typedef struct TF_ShapeAndTypeList TF_ShapeAndTypeList;
+
+// API for manipulating TF_ShapeAndTypeList objects.
+//
+TF_CAPI_EXPORT extern TF_ShapeAndTypeList* TF_NewShapeAndTypeList(
+    int num_shapes);
+TF_CAPI_EXPORT extern void TF_ShapeAndTypeListSetShape(
+    TF_ShapeAndTypeList* shape_list, int index, const int64_t* dims,
+    int num_dims);
+TF_CAPI_EXPORT extern void TF_ShapeAndTypeListSetUnknownShape(
+    TF_ShapeAndTypeList* shape_list, int index);
+TF_CAPI_EXPORT extern void TF_ShapeAndTypeListSetDtype(
+    TF_ShapeAndTypeList* shape_list, int index, TF_DataType dtype);
+TF_CAPI_EXPORT extern void TF_DeleteShapeAndTypeList(
+    TF_ShapeAndTypeList* shape_list);
+TF_CAPI_EXPORT extern void TF_DeleteShapeAndTypeListArray(
+    TF_ShapeAndTypeList** shape_list_array, int num_items);
+
+// Infer shapes for the given `op`. The arguments mimic the arguments of the
+// `shape_inference::InferenceContext` constructor. Note the following:
+//   - The inputs of the `op` are not used for shape inference. So, it is
+//     OK to not have the inputs properly set in `op`. See `input_tensors`
+//     if you want shape inference to consider the input tensors of the
+//     op for shape inference.
+//   - The types need not be set in `input_shapes` as it is not used.
+//   - The number of `input_tensors` should be the same as the number of items
+//     in `input_shapes`.
+//
+// The results are returned in `output_shapes` and
+// `output_resource_shapes_and_types`. The caller is responsible for freeing the
+// memory in these buffers by calling `TF_DeleteShapeAndTypeList`.
+TF_CAPI_EXPORT extern void TFE_InferShapes(
+    TFE_Op* op, TF_ShapeAndTypeList* input_shapes, TF_Tensor** input_tensors,
+    TF_ShapeAndTypeList* input_tensor_as_shapes,
+    TF_ShapeAndTypeList** input_resource_shapes_and_types,
+    TF_ShapeAndTypeList** output_shapes,
+    TF_ShapeAndTypeList*** output_resource_shapes_and_types, TF_Status* status);
+
+TF_CAPI_EXPORT extern void
+TF_ImportGraphDefOptionsSetValidateColocationConstraints(
+    TF_ImportGraphDefOptions* opts, unsigned char enable);
+
+// Load the library specified by library_filename and register the pluggable
+// device and related kernels present in that library. This function is not
+// supported on embedded on mobile and embedded platforms and will fail if
+// called.
+//
+// Pass "library_filename" to a platform-specific mechanism for dynamically
+// loading a library. The rules for determining the exact location of the
+// library are platform-specific and are not documented here.
+//
+// On success, returns the newly created library handle and places OK in status.
+// The caller owns the library handle.
+//
+// On failure, returns nullptr and places an error status in status.
+TF_CAPI_EXPORT extern TF_Library* TF_LoadPluggableDeviceLibrary(
+    const char* library_filename, TF_Status* status);
+
+// Frees the memory associated with the library handle.
+// Does NOT unload the library.
+TF_CAPI_EXPORT extern void TF_DeletePluggableDeviceLibraryHandle(
+    TF_Library* lib_handle);
+
+// Removes `func_name` from `g`. If `func_name` is not in `g`, an error will be
+// returned.
+TF_CAPI_EXPORT extern void TF_GraphRemoveFunction(TF_Graph* g,
+                                                  const char* func_name,
+                                                  TF_Status* status);
+
+#ifdef __cplusplus
+} /* end extern "C" */
+#endif
+
+#endif  // TENSORFLOW_C_C_API_EXPERIMENTAL_H_

tensorflow/c/c_api_experimental_test.cc

@@ -0,0 +1,365 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/c_api_experimental.h"
+
+#include "absl/types/optional.h"
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/c_test_util.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(CAPI_EXPERIMENTAL, GetServerDefTest) {
+  const std::string expected_text_proto(R"(cluster {
+  job {
+    name: "worker"
+    tasks {
+      key: 0
+      value: "tpuserver:0"
+    }
+    tasks {
+      key: 1
+      value: "localhost:1"
+    }
+  }
+}
+job_name: "worker"
+task_index: 1
+protocol: "grpc"
+)");
+
+  TF_Status* status = TF_NewStatus();
+  TF_Buffer* result = TFE_GetServerDef(expected_text_proto.c_str(), status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK);
+
+  ServerDef actual;
+  ASSERT_TRUE(actual.ParseFromArray(result->data, result->length));
+  std::string actual_text_proto;
+  tensorflow::protobuf::TextFormat::PrintToString(actual, &actual_text_proto);
+  EXPECT_EQ(expected_text_proto, actual_text_proto);
+
+  const std::string malformed_text_proto(R"(cluster {
+  job {
+    name: "worker")");
+  TF_Buffer* null_result =
+      TFE_GetServerDef(malformed_text_proto.c_str(), status);
+  EXPECT_NE(TF_GetCode(status), TF_OK);
+  EXPECT_TRUE(absl::StrContains(TF_Message(status),
+                                "Invalid text proto for ServerDef"));
+  EXPECT_EQ(null_result, nullptr);
+
+  // Cleanup
+  TF_DeleteBuffer(result);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI_EXPERIMENTAL, IsStateful) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  int assign = TF_OpIsStateful("AssignAddVariableOp", status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  EXPECT_EQ(assign, 1);
+  int id = TF_OpIsStateful("Identity", status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  EXPECT_EQ(id, 0);
+}
+
+class ShapeInferenceTest : public ::testing::Test {
+ protected:
+  ShapeInferenceTest()
+      : status_(TF_NewStatus()), tfe_context_options_(TFE_NewContextOptions()) {
+    tfe_context_ = TFE_NewContext(tfe_context_options_, status_);
+    CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+  }
+
+  ~ShapeInferenceTest() override {
+    TFE_DeleteContextOptions(tfe_context_options_);
+    TFE_DeleteContext(tfe_context_);
+    TF_DeleteStatus(status_);
+  }
+
+  // Checks the expected result of shape inference for the given `op`.
+  void CheckOutputShapes(
+      TFE_Op* op,
+      const std::vector<absl::optional<std::vector<int64_t>>>& input_shapes_vec,
+      const std::vector<TF_Tensor*>& input_tensors,
+      const absl::optional<std::vector<int64_t>>& expected_shape) {
+    // Create input_shapes.
+    TF_ShapeAndTypeList* input_shapes =
+        TF_NewShapeAndTypeList(input_shapes_vec.size());
+    for (size_t i = 0; i < input_shapes_vec.size(); ++i) {
+      const auto& input_shape = input_shapes_vec[i];
+      if (input_shape.has_value()) {
+        TF_ShapeAndTypeListSetShape(input_shapes, i, input_shape->data(),
+                                    input_shape->size());
+      } else {
+        TF_ShapeAndTypeListSetUnknownShape(input_shapes, i);
+      }
+    }
+    TF_ShapeAndTypeList* output_shapes;
+    TFE_InferShapes(op, input_shapes,
+                    input_tensors.empty()
+                        ? nullptr
+                        : const_cast<TF_Tensor**>(input_tensors.data()),
+                    /*input_tensors_as_shapes*/ nullptr,
+                    /*input_resource_shapes_and_types*/ nullptr, &output_shapes,
+                    /*output_resource_shapes_and_types*/ nullptr, status_);
+    CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+    CHECK_EQ(output_shapes->num_items, 1);
+
+    int num_dims = output_shapes->items[0].num_dims;
+    int64_t* dims = output_shapes->items[0].dims;
+
+    if (!expected_shape.has_value()) {
+      EXPECT_EQ(num_dims, -1);
+      EXPECT_EQ(dims, nullptr);
+      return;
+    }
+
+    EXPECT_EQ(num_dims, expected_shape->size());
+    for (size_t i = 0; i < num_dims; ++i) {
+      EXPECT_EQ(dims[i], (*expected_shape)[i]);
+    }
+    TF_DeleteShapeAndTypeList(input_shapes);
+    TF_DeleteShapeAndTypeList(output_shapes);
+  }
+
+  absl::optional<std::vector<int64_t>> make_shape(
+      std::vector<int64_t>&& dims) const {
+    return std::make_optional(dims);
+  }
+
+  absl::optional<std::vector<int64_t>> unknown_shape() const {
+    return std::nullopt;
+  }
+
+  static constexpr int64_t kUnknownDim =
+      shape_inference::InferenceContext::kUnknownDim;
+  TF_Status* status_;
+  TFE_ContextOptions* tfe_context_options_;
+  TFE_Context* tfe_context_;
+};
+
+TEST_F(ShapeInferenceTest, InfersShapesFromInputShapes) {
+  TFE_Op* matmul_op;
+  matmul_op = TFE_NewOp(tfe_context_, "MatMul", status_);
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+
+  // Infer shape when everything is known.
+  CheckOutputShapes(matmul_op,
+                    /*input_shapes*/ {make_shape({3, 2}), make_shape({2, 4})},
+                    /*input_tensors*/ {},
+                    /*expected_shape*/ make_shape({3, 4}));
+
+  // Infer shape when second operand has unknown shape.
+  CheckOutputShapes(matmul_op,
+                    /*input_shapes*/ {make_shape({3, 2}), unknown_shape()},
+                    /*input_tensors*/ {},
+                    /*expected_shape*/ make_shape({3, kUnknownDim}));
+
+  // Infer shape when some dimensions are unknown.
+  CheckOutputShapes(
+      matmul_op,
+      /*input_shapes*/ {make_shape({kUnknownDim, 2}), make_shape({2, 4})},
+      /*input_tensors*/ {},
+      /*expected_shape*/ make_shape({kUnknownDim, 4}));
+
+  // Infer shape when everything is unknown.
+  CheckOutputShapes(matmul_op,
+                    /*input_shapes*/ {unknown_shape(), unknown_shape()},
+                    /*input_tensors*/ {},
+                    /*expected_shape*/ make_shape({kUnknownDim, kUnknownDim}));
+
+  TFE_DeleteOp(matmul_op);
+  // TODO(bgogul): Add some death tests where status is not OK.
+}
+
+TEST_F(ShapeInferenceTest, InfersShapesFromInputTensors) {
+  // Prepare some tensors for shape.
+  TF_Tensor* tensor_1X6 = Int32Tensor({1, 6});
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+  TF_Tensor* tensor_1X1X6 = Int32Tensor({1, 1, 6});
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+
+  TFE_Op* reshape_op = TFE_NewOp(tfe_context_, "Reshape", status_);
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+  TFE_OpSetAttrType(reshape_op, "T", TF_FLOAT);
+  TFE_OpSetAttrType(reshape_op, "Tshape", TF_INT32);
+  CheckOutputShapes(reshape_op,
+                    /* input_shapes*/ {unknown_shape(), unknown_shape()},
+                    /* input_tensors*/ {nullptr, tensor_1X6},
+                    /*expected_shape*/ make_shape({1, 6}));
+  TFE_DeleteOp(reshape_op);
+  reshape_op = nullptr;
+
+  TFE_Op* fill_op = TFE_NewOp(tfe_context_, "Fill", status_);
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+  TFE_OpSetAttrType(fill_op, "T", TF_FLOAT);
+  TFE_OpSetAttrType(fill_op, "Tshape", TF_INT32);
+
+  float five = 5.0;
+  TFE_TensorHandle* scalar = TestScalarTensorHandle(tfe_context_, five);
+  TF_Tensor* scalarTensor = TFE_TensorHandleResolve(scalar, status_);
+  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
+  CheckOutputShapes(fill_op,
+                    /* input_shapes*/ {unknown_shape(), unknown_shape()},
+                    /* input_tensors*/ {tensor_1X1X6, scalarTensor},
+                    /*expected_shape*/ make_shape({1, 1, 6}));
+  TFE_DeleteOp(fill_op);
+  fill_op = nullptr;
+
+  TFE_DeleteTensorHandle(scalar);
+  TF_DeleteTensor(scalarTensor);
+  TF_DeleteTensor(tensor_1X1X6);
+  TF_DeleteTensor(tensor_1X6);
+}
+
+TEST(CAPI_EXPERIMENTAL, LibraryPluggableDeviceLoadFunctions) {
+  // TODO(penpornk): Enable this test on Windows.
+#if !defined(PLATFORM_WINDOWS)
+#if !defined(TENSORFLOW_NO_SHARED_OBJECTS)
+  // Load the library.
+  TF_Status* status = TF_NewStatus();
+  string lib_path =
+      tensorflow::GetDataDependencyFilepath(tensorflow::io::JoinPath(
+          "tensorflow", "c", "experimental", "stream_executor", "test",
+          "test_pluggable_device.so"));
+  TF_Library* lib = TF_LoadPluggableDeviceLibrary(lib_path.c_str(), status);
+  TF_Code code = TF_GetCode(status);
+  string status_msg(TF_Message(status));
+  TF_DeleteStatus(status);
+  ASSERT_EQ(TF_OK, code) << status_msg;
+  TF_DeletePluggableDeviceLibraryHandle(lib);
+#endif  // !defined(TENSORFLOW_NO_SHARED_OBJECTS)
+#endif  // !defined(PLATFORM_WINDOWS)
+}
+
+TEST(CAPI_EXPERIMENTAL, LibraryNextPluggableDeviceLoadFunctions) {
+  // TODO(penpornk): Enable this test on Windows.
+#if !defined(PLATFORM_WINDOWS)
+#if !defined(TENSORFLOW_NO_SHARED_OBJECTS)
+  // Load the library.
+  TF_Status* status = TF_NewStatus();
+  string lib_path =
+      tensorflow::GetDataDependencyFilepath(tensorflow::io::JoinPath(
+          "tensorflow", "core", "common_runtime", "next_pluggable_device", "c",
+          "test_next_pluggable_device_plugin.so"));
+  TF_Library* lib = TF_LoadPluggableDeviceLibrary(lib_path.c_str(), status);
+  TF_Code code = TF_GetCode(status);
+  string status_msg(TF_Message(status));
+  TF_DeleteStatus(status);
+  ASSERT_EQ(TF_OK, code) << status_msg;
+  TF_DeletePluggableDeviceLibraryHandle(lib);
+#endif  // !defined(TENSORFLOW_NO_SHARED_OBJECTS)
+#endif  // !defined(PLATFORM_WINDOWS)
+}
+
+void DefineFunction(const char* name, TF_Function** func,
+                    const char* description = nullptr,
+                    bool append_hash = false) {
+  std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)> func_graph(
+      TF_NewGraph(), TF_DeleteGraph);
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> s(TF_NewStatus(),
+                                                           TF_DeleteStatus);
+
+  TF_Operation* feed = Placeholder(func_graph.get(), s.get());
+  TF_Operation* neg = Neg(feed, func_graph.get(), s.get());
+
+  TF_Output inputs[] = {{feed, 0}};
+  TF_Output outputs[] = {{neg, 0}};
+  *func = TF_GraphToFunction(func_graph.get(), name, append_hash, -1,
+                             /*opers=*/nullptr, 1, inputs, 1, outputs,
+                             /*output_names=*/nullptr,
+                             /*opts=*/nullptr, description, s.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(s.get())) << TF_Message(s.get());
+  ASSERT_NE(*func, nullptr);
+}
+
+class CApiExperimentalFunctionTest : public ::testing::Test {
+ protected:
+  CApiExperimentalFunctionTest()
+      : s_(TF_NewStatus()), func_graph_(TF_NewGraph()), func_(nullptr) {}
+
+  void SetUp() override {}
+
+  ~CApiExperimentalFunctionTest() override {
+    TF_DeleteFunction(func_);
+    TF_DeleteGraph(func_graph_);
+    TF_DeleteStatus(s_);
+  }
+
+  const char* func_name_ = "MyFunc";
+  TF_Status* s_;
+  TF_Graph* func_graph_;
+  TF_Function* func_;
+};
+
+TEST_F(CApiExperimentalFunctionTest, GraphRemoveFunction) {
+  TF_Function* funcs[1];
+  DefineFunction(func_name_, &func_);
+
+  TF_GraphCopyFunction(func_graph_, func_, nullptr, s_);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  EXPECT_EQ(TF_GraphNumFunctions(func_graph_), 1);
+  EXPECT_EQ(TF_GraphGetFunctions(func_graph_, funcs, 1, s_), 1);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  TF_GraphRemoveFunction(func_graph_, func_name_, s_);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  EXPECT_EQ(TF_GraphNumFunctions(func_graph_), 0);
+  EXPECT_EQ(TF_GraphGetFunctions(func_graph_, funcs, 1, s_), 0);
+
+  TF_DeleteFunction(funcs[0]);
+}
+
+TEST_F(CApiExperimentalFunctionTest, EmptyGraphRemoveNonExistentFunction) {
+  TF_GraphRemoveFunction(func_graph_, "wrong_name", s_);
+  EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_));
+  EXPECT_EQ(std::string("Tried to remove non-existent function 'wrong_name'."),
+            std::string(TF_Message(s_)));
+}
+
+TEST_F(CApiExperimentalFunctionTest, GraphRemoveNonExistentFunction) {
+  TF_Function* funcs[1];
+  DefineFunction(func_name_, &func_);
+
+  TF_GraphCopyFunction(func_graph_, func_, nullptr, s_);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  EXPECT_EQ(TF_GraphNumFunctions(func_graph_), 1);
+  EXPECT_EQ(TF_GraphGetFunctions(func_graph_, funcs, 1, s_), 1);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  TF_GraphRemoveFunction(func_graph_, "wrong_name", s_);
+  EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_));
+  EXPECT_EQ(std::string("Tried to remove non-existent function 'wrong_name'."),
+            std::string(TF_Message(s_)));
+  TF_DeleteFunction(funcs[0]);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/c_api_function.cc

@@ -0,0 +1,366 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+
+#include "absl/strings/match.h"
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/tf_buffer_internal.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor.pb.h"  // NOLINT
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/platform/base64.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/util/debug_data_dumper.h"
+
+using tensorflow::errors::InvalidArgument;
+
+namespace tensorflow {
+namespace {
+
+absl::Status ValidateNonRefOutput(const Node* node, int idx) {
+  const DataType& dt = node->output_type(idx);
+  return IsRefType(dt) ? absl::InvalidArgumentError(absl::StrCat(
+                             "Output ", idx, " of node '", node->name(),
+                             "' has a reference type ", DataTypeString(dt)))
+                       : absl::OkStatus();
+}
+
+// Converts `ninputs` and `inputs` into `inputs_tensors` and `input_nodes` and
+// does various checks while doing so. `input_nodes` will contain the same
+// information as input_tensors just in a different structure to make
+// following processing easier. TODO(iga): Simplify this nested structure.
+absl::Status ProcessInputs(
+    const TF_Graph* fn_body, const char* fn_name, int ninputs,
+    const TF_Output* inputs, std::vector<OutputTensor>* input_tensors,
+    std::unordered_map<const Node*, std::vector<int>>* input_nodes)
+    TF_EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) {
+  input_tensors->reserve(ninputs);
+  for (int i = 0; i < ninputs; ++i) {
+    Node* node = inputs[i].oper ? &inputs[i].oper->node : nullptr;
+    int idx = inputs[i].index;
+
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        fn_body->graph.IsValidOutputTensor(node, idx),
+        "Encountered while processing input ", i, " into function '", fn_name,
+        "'");
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(node, idx),
+                                    "Encountered while processing input ", i,
+                                    " into function '", fn_name, "'");
+
+    input_tensors->emplace_back(node, idx);
+
+    const auto& iter = input_nodes->find(node);
+    if (iter == input_nodes->end()) {
+      input_nodes->insert({node, {idx}});
+    } else {
+      auto& indices = iter->second;
+      if (std::find(indices.begin(), indices.end(), idx) != indices.end()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("TF_Output ", node->name(), ":", idx,
+                         " appears more than once in the input list"));
+      }
+      indices.push_back(idx);
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Converts `noutputs` and `outputs` into `outputs_tensors` and does various
+// checks while doing so.
+absl::Status ProcessOutputs(const TF_Graph* fn_body, const char* fn_name,
+                            int noutputs, const TF_Output* outputs,
+                            std::vector<OutputTensor>* output_tensors)
+    TF_EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) {
+  output_tensors->reserve(noutputs);
+  for (int i = 0; i < noutputs; ++i) {
+    Node* node = outputs[i].oper ? &outputs[i].oper->node : nullptr;
+    int idx = outputs[i].index;
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        fn_body->graph.IsValidOutputTensor(node, idx),
+        "Encountered while processing output ", i, " from function '", fn_name,
+        "'");
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(node, idx),
+                                    "Encountered while creating function '",
+                                    fn_name, "'");
+    output_tensors->emplace_back(node, idx);
+  }
+  return absl::OkStatus();
+}
+
+// Populates `body_nodes` with the nodes that will become function's body.
+// Performs various checks.
+absl::Status ComputeBodyNodes(
+    const TF_Graph* fn_body, const char* fn_name, int num_opers,
+    const TF_Operation* const* opers,
+    const std::unordered_map<const Node*, std::vector<int>>& input_nodes,
+    std::vector<const Node*>* body_nodes)
+    TF_EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) {
+  if (num_opers == -1) {
+    for (const Node* node : fn_body->graph.op_nodes()) {
+      const auto& iter = input_nodes.find(node);
+      if (iter == input_nodes.end()) {
+        // This node is not referenced in inputs. Add it to the body.
+        body_nodes->push_back(node);
+      } else {
+        // This node is referenced in inputs. Currently, we place an
+        // artificial restriction and require that when num_opers=-1, such
+        // nodes must have a single output.
+        if (node->num_outputs() != 1) {
+          return absl::InvalidArgumentError(absl::StrCat(
+              "When `num_opers` is set to -1, nodes referenced in `inputs` "
+              "must have a single output. Node ",
+              node->name(), " has ", node->num_outputs(),
+              " outputs. Encountered while creating function '", fn_name, "'"));
+        }
+      }
+    }
+  } else {
+    body_nodes->reserve(num_opers);
+    for (int i = 0; i < num_opers; ++i) {
+      const Node* node = &opers[i]->node;
+      body_nodes->push_back(node);
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+}  // namespace tensorflow
+
+using tensorflow::Node;
+using tensorflow::string;
+
+TF_Function* TF_GraphToFunctionWithControlOutputs(
+    const TF_Graph* fn_body, const char* fn_name,
+    unsigned char append_hash_to_fn_name, int num_opers,
+    const TF_Operation* const* opers, int ninputs, const TF_Output* inputs,
+    int noutputs, const TF_Output* outputs, const char* const* output_names,
+    int ncontrol_outputs, const TF_Operation* const* control_outputs,
+    const char* const* control_output_names, const TF_FunctionOptions* opts,
+    const char* description, TF_Status* status) {
+  tensorflow::mutex_lock l(fn_body->mu);
+
+  // Process inputs.
+  std::vector<tensorflow::OutputTensor> input_tensors;
+  std::unordered_map<const Node*, std::vector<int>> input_nodes;
+  status->status = tensorflow::ProcessInputs(fn_body, fn_name, ninputs, inputs,
+                                             &input_tensors, &input_nodes);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  // Process outputs.
+  std::vector<tensorflow::OutputTensor> output_tensors;
+  status->status = tensorflow::ProcessOutputs(fn_body, fn_name, noutputs,
+                                              outputs, &output_tensors);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  // Process output names.
+  std::vector<string> output_names_vec;
+  if (output_names) {
+    output_names_vec.reserve(noutputs);
+    for (int i = 0; i < noutputs; ++i) {
+      output_names_vec.push_back(string(output_names[i]));
+    }
+  }
+
+  // Process control output names.
+  std::vector<string> control_output_names_vec;
+  if (control_output_names) {
+    control_output_names_vec.reserve(ncontrol_outputs);
+    for (int i = 0; i < ncontrol_outputs; ++i) {
+      control_output_names_vec.push_back(string(control_output_names[i]));
+    }
+  }
+
+  // Compute body nodes.
+  std::vector<const Node*> body_nodes;
+  status->status = tensorflow::ComputeBodyNodes(
+      fn_body, fn_name, num_opers, opers, input_nodes, &body_nodes);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  // Compute body nodes.
+  std::vector<const Node*> control_output_nodes;
+  control_output_nodes.reserve(ncontrol_outputs);
+  for (int i = 0; i < ncontrol_outputs; ++i) {
+    control_output_nodes.push_back(&control_outputs[i]->node);
+  }
+
+  // Do the actual function creation.
+  DCHECK(append_hash_to_fn_name <= 1);
+  tensorflow::FunctionDef fdef;
+  status->status = tensorflow::GraphToFunctionDef(
+      fn_body->graph, fn_name, append_hash_to_fn_name != 0,
+      /*set_stateful_from_nodes=*/true,
+      /*copy_placeholder_attrs_from_nodes=*/true, body_nodes, input_tensors,
+      output_tensors, output_names_vec, control_output_nodes,
+      control_output_names_vec, description, &fdef);
+  if (TF_GetCode(status) != TF_OK) {
+    return nullptr;
+  }
+
+  // Dump the op creation stacktraces for debugging purpose.
+  DEBUG_DATA_DUMPER()->DumpOpCreationStackTraces(
+      fn_name, kDebugGroupOpStacktrace, "initial", &fn_body->graph);
+
+  tensorflow::StackTracesMap stack_traces;
+  for (const Node* n : fn_body->graph.nodes()) {
+    stack_traces[n->name()] = n->GetStackTrace();
+  }
+
+  TF_Function* tf_function = new TF_Function();
+  tf_function->record = new tensorflow::FunctionRecord(
+      std::move(fdef), std::move(stack_traces), false);
+
+  return tf_function;
+}
+
+TF_Function* TF_GraphToFunction(const TF_Graph* fn_body, const char* fn_name,
+                                unsigned char append_hash_to_fn_name,
+                                int num_opers, const TF_Operation* const* opers,
+                                int ninputs, const TF_Output* inputs,
+                                int noutputs, const TF_Output* outputs,
+                                const char* const* output_names,
+                                const TF_FunctionOptions* opts,
+                                const char* description, TF_Status* status) {
+  return TF_GraphToFunctionWithControlOutputs(
+      fn_body, fn_name, append_hash_to_fn_name, num_opers, opers, ninputs,
+      inputs, noutputs, outputs, output_names, 0, nullptr, nullptr, opts,
+      description, status);
+}
+
+const char* TF_FunctionName(TF_Function* func) {
+  return func->record->fdef().signature().name().c_str();
+}
+
+void TF_GraphCopyFunction(TF_Graph* g, const TF_Function* func,
+                          const TF_Function* grad, TF_Status* status) {
+  if (func == nullptr) {
+    status->status = absl::InvalidArgumentError(
+        "'func' argument to TF_GraphCopyFunction cannot be null");
+    return;
+  }
+
+  tensorflow::mutex_lock l(g->mu);
+  status->status = g->graph.AddFunctionDef(func->record->fdef(),
+                                           func->record->stack_traces());
+  if (TF_GetCode(status) != TF_OK) return;
+  if (!grad) return;
+
+  status->status = g->graph.AddFunctionDef(grad->record->fdef(),
+                                           grad->record->stack_traces());
+  if (TF_GetCode(status) != TF_OK) return;
+
+  tensorflow::GradientDef gdef;
+  gdef.set_function_name(func->record->fdef().signature().name());
+  gdef.set_gradient_func(grad->record->fdef().signature().name());
+  status->status = g->graph.AddGradientDef(std::move(gdef));
+}
+
+int TF_GraphNumFunctions(TF_Graph* g) {
+  tensorflow::mutex_lock l(g->mu);
+  return g->graph.flib_def().num_functions();
+}
+
+int TF_GraphGetFunctions(TF_Graph* g, TF_Function** funcs, int max_func,
+                         TF_Status* status) {
+  tensorflow::FunctionDefLibrary lib;
+  {
+    tensorflow::mutex_lock l(g->mu);
+    lib = g->graph.flib_def().ToProto();
+  }
+  const auto len = std::min(max_func, static_cast<int>(lib.function_size()));
+  for (int i = 0; i < len; ++i) {
+    TF_Function* func = new TF_Function();
+    func->record = new tensorflow::FunctionRecord(lib.function(i), {}, false);
+    funcs[i] = func;
+  }
+  status->status = absl::OkStatus();
+  return len;
+}
+
+void TF_FunctionToFunctionDef(TF_Function* func, TF_Buffer* output_func_def,
+                              TF_Status* status) {
+  status->status = MessageToBuffer(func->record->fdef(), output_func_def);
+}
+
+TF_Function* TF_FunctionImportFunctionDef(const void* proto, size_t proto_len,
+                                          TF_Status* status) {
+  tensorflow::FunctionDef fdef;
+  bool success = fdef.ParseFromArray(proto, proto_len);
+  if (!success) {
+    status->status = absl::InvalidArgumentError(
+        "Invalid FunctionDef given to TF_FunctionImportFunctionDef");
+    return nullptr;
+  }
+
+  TF_Function* func = new TF_Function();
+  func->record = new tensorflow::FunctionRecord(std::move(fdef), {}, false);
+  status->status = absl::OkStatus();
+  return func;
+}
+
+void TF_FunctionSetAttrValueProto(TF_Function* func, const char* attr_name,
+                                  const void* proto, size_t proto_len,
+                                  TF_Status* status) {
+  tensorflow::AttrValue attr_value;
+  if (!attr_value.ParseFromArray(proto, proto_len)) {
+    status->status = absl::InvalidArgumentError(
+        "Unparseable AttrValue proto passed to "
+        "TF_FunctionSetAttrValueProto");
+    return;
+  }
+
+  auto fdef_or = func->record->mutable_fdef();
+  if (!fdef_or.ok()) {
+    status->status = fdef_or.status();
+    return;
+  }
+
+  (*(fdef_or.value()->mutable_attr()))[string(attr_name)] = attr_value;
+
+  status->status = absl::OkStatus();
+}
+
+void TF_FunctionGetAttrValueProto(TF_Function* func, const char* attr_name,
+                                  TF_Buffer* output_attr_value,
+                                  TF_Status* status) {
+  const auto& it = func->record->fdef().attr().find(attr_name);
+  if (it == func->record->fdef().attr().end()) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("Function '", func->record->fdef().signature().name(),
+                     "' has no attr named '", attr_name, "'."));
+    return;
+  }
+  status->status = MessageToBuffer(it->second, output_attr_value);
+}
+
+void TF_DeleteFunction(TF_Function* func) {
+  if (func == nullptr) {
+    return;
+  }
+
+  func->record->Unref();
+  func->record = nullptr;
+  delete func;
+}

tensorflow/c/c_api_internal.h

@@ -0,0 +1,226 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_API_INTERNAL_H_
+#define TENSORFLOW_C_C_API_INTERNAL_H_
+
+#include <deque>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/c/c_api.h"
+
+// clang-format off
+// Required for IS_MOBILE_PLATFORM
+#include "tensorflow/core/platform/platform.h"
+// clang-format on
+
+#include "tensorflow/c/tf_status_internal.h"
+#include "tensorflow/c/tf_tensor_internal.h"
+#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+#include "tensorflow/core/framework/op_gen_lib.h"
+#endif  // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/common_runtime/shape_refiner.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+class Device;
+class DeviceMgr;
+class ServerInterface;
+}  // namespace tensorflow
+
+// Internal structures used by the C API. These are likely to change and should
+// not be depended on.
+
+struct TF_SessionOptions {
+  tensorflow::SessionOptions options;
+};
+
+struct TF_DeprecatedSession {
+  tensorflow::Session* session;
+};
+
+struct TF_Library {
+  void* lib_handle;
+  TF_Buffer op_list;
+};
+
+struct TF_Graph {
+  TF_Graph();
+
+  mutable tensorflow::mutex mu;
+  tensorflow::Graph graph TF_GUARDED_BY(mu);
+
+  // Runs shape inference.
+  tensorflow::ShapeRefiner refiner TF_GUARDED_BY(mu);
+
+  // Maps from name of an operation to the Node* in 'graph'.
+  std::unordered_map<std::string, tensorflow::Node*> name_map TF_GUARDED_BY(mu);
+
+  // The keys of this map are all the active sessions using this graph. Each
+  // value records whether the graph has been mutated since the corresponding
+  // session has been run (this is detected in RecordMutation function). If the
+  // string is empty, no mutation has occurred. Otherwise the string is a
+  // description of the mutation suitable for returning to the user.
+  //
+  // Sessions are added to this map in TF_NewSession, and removed in
+  // TF_DeleteSession.
+  // TF_Graph may only / must be deleted when
+  //   sessions.size() == 0 && delete_requested == true
+  //
+  // TODO(b/74949947): mutations currently trigger a warning instead of a bad
+  // status, this should be reverted when possible.
+  tensorflow::gtl::FlatMap<TF_Session*, std::string> sessions TF_GUARDED_BY(mu);
+  bool delete_requested TF_GUARDED_BY(mu);  // set true by TF_DeleteGraph
+
+  // Used to link graphs contained in TF_WhileParams to the parent graph that
+  // will eventually contain the full while loop.
+  TF_Graph* parent;
+  TF_Output* parent_inputs;
+};
+
+struct TF_OperationDescription {
+  TF_OperationDescription(TF_Graph* g, const char* op_type,
+                          const char* node_name)
+      : node_builder(node_name, op_type, g->graph.op_registry()), graph(g) {}
+
+  tensorflow::NodeBuilder node_builder;
+  TF_Graph* graph;
+  std::set<std::string> colocation_constraints;
+};
+
+struct TF_Operation {
+  tensorflow::Node node;
+
+ private:
+  ~TF_Operation() = default;
+};
+
+struct TF_Session {
+  TF_Session(tensorflow::Session* s, TF_Graph* g);
+
+  tensorflow::Session* session;
+  TF_Graph* const graph;
+
+  tensorflow::mutex mu TF_ACQUIRED_AFTER(TF_Graph::mu);
+  int last_num_graph_nodes;
+
+  // If true, TF_SessionRun and similar methods will call
+  // ExtendSessionGraphHelper before running the graph (this is the default
+  // public behavior). Can be set to false if the caller needs to call
+  // ExtendSessionGraphHelper manually.
+  std::atomic<bool> extend_before_run;
+};
+
+struct TF_ImportGraphDefOptions {
+  tensorflow::ImportGraphDefOptions opts;
+
+  // Backing memory for TensorId fields in opts.
+  // TODO(skyewm): it'd be better if ImportGraphDefOptions owned this.
+  // A deque is necessary to preserve pointer stability in the code.
+  std::deque<std::string> tensor_id_data;
+};
+
+struct TF_ImportGraphDefResults {
+  std::vector<TF_Output> return_tensors;
+  std::vector<TF_Operation*> return_nodes;
+  std::vector<const char*> missing_unused_key_names;
+  std::vector<int> missing_unused_key_indexes;
+
+  // Backing memory for missing_unused_key_names values.
+  std::deque<std::string> missing_unused_key_names_data;
+};
+
+struct TF_DeviceList {
+  std::vector<tensorflow::DeviceAttributes> response;
+};
+
+struct TF_Function {
+  tensorflow::FunctionRecord* record;
+};
+
+struct TF_ApiDefMap {
+  explicit TF_ApiDefMap(const tensorflow::OpList& op_list)
+      :
+#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+        api_def_map(op_list),
+#endif  // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+        update_docs_called(false) {
+  }
+
+#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+  tensorflow::ApiDefMap api_def_map TF_GUARDED_BY(lock);
+#endif  // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+  bool update_docs_called TF_GUARDED_BY(lock);
+  tensorflow::mutex lock;
+};
+
+#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+struct TF_Server {
+  TF_Server(std::unique_ptr<tensorflow::ServerInterface> server);
+
+  const std::string target;
+  std::unique_ptr<tensorflow::ServerInterface> server;
+};
+#endif  // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)
+
+namespace tensorflow {
+
+// Set the shapes and types of the output's handle.
+//
+// The lengths of the arrays pointed to by `shapes`, `ranks`, and `types` must
+// all be equal to `num_shapes_and_types`. If `ranks[i] != -1`, (i.e., if the
+// rank is known), then it must be equal to the length of `shapes[i]`; if
+// `ranks[i] == 1`, then `shapes[i]` may be nullptr.
+//
+// TODO(akshayka): Implement a corresponding getter method.
+void TF_GraphSetOutputHandleShapesAndTypes(TF_Graph* graph, TF_Output output,
+                                           int num_shapes_and_types,
+                                           const int64_t** shapes,
+                                           const int* ranks,
+                                           const TF_DataType* types,
+                                           TF_Status* status);
+
+void RecordMutation(TF_Graph* graph, const TF_Operation& op,
+                    const char* mutation_type)
+    TF_EXCLUSIVE_LOCKS_REQUIRED(graph->mu);
+
+bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status)
+    TF_LOCKS_EXCLUDED(session->graph->mu, session->mu);
+
+std::string getTF_OutputDebugString(TF_Output node);
+
+// Set whether to propagate assigned device information when constructing a new
+// Graph from a GraphDef. By default assigned device information is not copied
+// and is re-computed by the runtime.
+inline void TF_ImportGraphDefOptionsSetPropagateDeviceSpec(
+    TF_ImportGraphDefOptions* opts, unsigned char propagate_device_spec) {
+  opts->opts.propagate_device_spec = propagate_device_spec;
+}
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_C_C_API_INTERNAL_H_

tensorflow/c/c_api_macros.h

@@ -0,0 +1,51 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_API_MACROS_H_
+#define TENSORFLOW_C_C_API_MACROS_H_
+
+#ifdef SWIG
+#define TF_CAPI_EXPORT
+#else
+#if defined(_WIN32)
+#ifdef TF_COMPILE_LIBRARY
+#define TF_CAPI_EXPORT __declspec(dllexport)
+#else
+#define TF_CAPI_EXPORT __declspec(dllimport)
+#endif  // TF_COMPILE_LIBRARY
+#else
+#ifdef TF_CAPI_WEAK
+#define TF_CAPI_EXPORT \
+  __attribute__((visibility("default"))) __attribute((weak))
+#else
+#define TF_CAPI_EXPORT __attribute__((visibility("default")))
+#endif  // TF_CAPI_WEAK
+#endif  // _WIN32
+#endif  // SWIG
+
+// TF_Bool is the C API typedef for unsigned char, while TF_BOOL is
+// the datatype for boolean tensors.
+#ifndef TF_Bool
+#define TF_Bool unsigned char
+#endif  // TF_Bool
+
+// Macro used to calculate struct size for maintaining ABI stability across
+// different struct implementations.
+#ifndef TF_OFFSET_OF_END
+#define TF_OFFSET_OF_END(TYPE, MEMBER) \
+  (offsetof(TYPE, MEMBER) + sizeof(((TYPE *)0)->MEMBER))
+#endif  // TF_OFFSET_OF_END
+
+#endif  // TENSORFLOW_C_C_API_MACROS_H_

tensorflow/c/c_api_macros_internal.h

@@ -0,0 +1,48 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_API_MACROS_INTERNAL_H_
+#define TENSORFLOW_C_C_API_MACROS_INTERNAL_H_
+
+#ifdef __cplusplus
+#include "tensorflow/core/platform/status.h"
+
+// Macro to verify that the field `struct_size` of STRUCT_OBJ is initialized.
+// `struct_size` is used for struct member compatibility check between core TF
+// and plug-ins with the same C API minor version. More info here:
+// https://github.com/tensorflow/community/blob/master/rfcs/20200612-stream-executor-c-api/C_API_versioning_strategy.md
+#define TF_VALIDATE_STRUCT_SIZE(STRUCT_NAME, STRUCT_OBJ, SIZE_VALUE_NAME) \
+  do {                                                                    \
+    if (STRUCT_OBJ.struct_size == 0) {                                    \
+      return tensorflow::Status(absl::StatusCode::kFailedPrecondition,    \
+                                "Expected initialized `" #STRUCT_NAME     \
+                                "` structure with `struct_size` field "   \
+                                "set to " #SIZE_VALUE_NAME                \
+                                ". Found `struct_size` = 0.");            \
+    }                                                                     \
+  } while (0)
+
+// Macro to verify that the field NAME of STRUCT_OBJ is not null.
+#define TF_VALIDATE_NOT_NULL(STRUCT_NAME, STRUCT_OBJ, NAME)            \
+  do {                                                                 \
+    if (STRUCT_OBJ.NAME == 0) {                                        \
+      return tensorflow::Status(absl::StatusCode::kFailedPrecondition, \
+                                "'" #NAME "' field in " #STRUCT_NAME   \
+                                " must be set.");                      \
+    }                                                                  \
+  } while (0)
+
+#endif  // __cplusplus
+#endif  // TENSORFLOW_C_C_API_MACROS_INTERNAL_H_

tensorflow/c/c_op_requires.h

@@ -0,0 +1,51 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_OP_REQUIRES_H_
+#define TENSORFLOW_C_C_OP_REQUIRES_H_
+
+#include "tensorflow/core/platform/macros.h"
+
+namespace tensorflow {
+
+// Convenience macros for asserting and handling exceptional conditions, for
+// C structs, including `TF_OpKernelContext`, `TF_Status`, etc. This is analogus
+// to the macros in tensorflow/core/framework/op_requires.h. This is provided
+// for plugin OpKernel developer's convenience.
+
+#define C_OPKERNELCONTEXT_REQUIRES_OK(CTX, C_STATUS, __VA_ARGS__) \
+  do {                                                            \
+    ::tensorflow::Status _s(__VA_ARGS__);                         \
+    if (!TF_PREDICT_TRUE(_s.ok())) {                              \
+      ::tensorflow::Set_TF_Status_from_Status(C_STATUS, _s);      \
+      TF_OpKernelContext_Failure(CTX, C_STATUS);                  \
+      TF_DeleteStatus(C_STATUS);                                  \
+      return;                                                     \
+    }                                                             \
+  } while (0)
+
+#define TF_CLEANUP_AND_RETURN_IF_ERROR(C_STATUS, BUFFER, __VA_ARGS__) \
+  do {                                                                \
+    ::tensorflow::Status _s(__VA_ARGS__);                             \
+    if (TF_PREDICT_TRUE(!_s.ok())) {                                  \
+      TF_DeleteStatus(C_STATUS);                                      \
+      TF_DeleteBuffer(BUFFER);                                        \
+      return _s;                                                      \
+    }                                                                 \
+  } while (0)
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_C_OP_REQUIRES_H_

tensorflow/c/c_test.c

@@ -0,0 +1,93 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <limits.h>
+#include <memory.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <unistd.h>
+
+#ifdef _WIN32
+#include <process.h>
+#endif
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/c/env.h"
+#include "tensorflow/c/kernels.h"
+
+// A create function. This will never actually get called in this test, it's
+// just nice to know that it compiles.
+void* create(TF_OpKernelConstruction* ctx) {
+  TF_DataType type;
+  TF_Status* s = TF_NewStatus();
+  TF_OpKernelConstruction_GetAttrType(ctx, "foobar", &type, s);
+  TF_DeleteStatus(s);
+  return NULL;
+}
+
+// A compute function. This will never actually get called in this test, it's
+// just nice to know that it compiles.
+void compute(void* kernel, TF_OpKernelContext* ctx) {
+  TF_Tensor* input;
+  TF_Status* s = TF_NewStatus();
+  TF_GetInput(ctx, 0, &input, s);
+  TF_DeleteTensor(input);
+  TF_DeleteStatus(s);
+}
+
+// Exercises tensorflow's C API.
+int main(int argc, char** argv) {
+  TF_InitMain(argv[0], &argc, &argv);
+
+  struct TF_StringStream* s = TF_GetLocalTempDirectories();
+  const char* path;
+
+  if (!TF_StringStreamNext(s, &path)) {
+    fprintf(stderr, "TF_GetLocalTempDirectories returned no results\n");
+    return 1;
+  }
+
+  char file_name[100];
+  time_t t = time(NULL);
+  snprintf(file_name, sizeof(file_name), "test-%d-%ld.txt", getpid(), t);
+
+  size_t length = 2 + strlen(path) + strlen(file_name);
+  char* full_path = malloc(length);
+  snprintf(full_path, length, "%s/%s", path, file_name);
+
+  TF_WritableFileHandle* h;
+  TF_Status* status = TF_NewStatus();
+  TF_NewWritableFile(full_path, &h, status);
+  if (TF_GetCode(status) != TF_OK) {
+    fprintf(stderr, "TF_NewWritableFile failed: %s\n", TF_Message(status));
+    return 1;
+  }
+  fprintf(stderr, "wrote %s\n", full_path);
+  free(full_path);
+  TF_CloseWritableFile(h, status);
+  if (TF_GetCode(status) != TF_OK) {
+    fprintf(stderr, "TF_CloseWritableFile failed: %s\n", TF_Message(status));
+  }
+  TF_StringStreamDone(s);
+
+  TF_KernelBuilder* b =
+      TF_NewKernelBuilder("SomeOp", "SomeDevice", &create, &compute, NULL);
+  TF_RegisterKernelBuilder("someKernel", b, status);
+
+  TF_DeleteStatus(status);
+  return 0;
+}

tensorflow/c/c_test_util.cc

@@ -0,0 +1,551 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/c_test_util.h"
+
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/public/session_options.h"
+
+using tensorflow::GraphDef;
+using tensorflow::NodeDef;
+
+static void BoolDeallocator(void* data, size_t, void* arg) {
+  delete[] static_cast<bool*>(data);
+}
+
+static void Int32Deallocator(void* data, size_t, void* arg) {
+  delete[] static_cast<int32_t*>(data);
+}
+
+static void DoubleDeallocator(void* data, size_t, void* arg) {
+  delete[] static_cast<double*>(data);
+}
+
+static void FloatDeallocator(void* data, size_t, void* arg) {
+  delete[] static_cast<float*>(data);
+}
+
+TF_Tensor* BoolTensor(bool v) {
+  const int num_bytes = sizeof(bool);
+  bool* values = new bool[1];
+  values[0] = v;
+  return TF_NewTensor(TF_BOOL, nullptr, 0, values, num_bytes, &BoolDeallocator,
+                      nullptr);
+}
+
+TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values) {
+  int64_t num_values = 1;
+  for (int i = 0; i < num_dims; ++i) {
+    num_values *= dims[i];
+  }
+  TF_Tensor* t =
+      TF_AllocateTensor(TF_INT8, dims, num_dims, sizeof(char) * num_values);
+  memcpy(TF_TensorData(t), values, sizeof(char) * num_values);
+  return t;
+}
+
+TF_Tensor* Int32Tensor(const int64_t* dims, int num_dims,
+                       const int32_t* values) {
+  int64_t num_values = 1;
+  for (int i = 0; i < num_dims; ++i) {
+    num_values *= dims[i];
+  }
+  TF_Tensor* t =
+      TF_AllocateTensor(TF_INT32, dims, num_dims, sizeof(int32_t) * num_values);
+  memcpy(TF_TensorData(t), values, sizeof(int32_t) * num_values);
+  return t;
+}
+
+TF_Tensor* Int32Tensor(const std::vector<int32_t>& values) {
+  int64_t dims = values.size();
+  return Int32Tensor(&dims, 1, values.data());
+}
+
+TF_Tensor* Int32Tensor(int32_t v) {
+  const int num_bytes = sizeof(int32_t);
+  int32_t* values = new int32_t[1];
+  values[0] = v;
+  return TF_NewTensor(TF_INT32, nullptr, 0, values, num_bytes,
+                      &Int32Deallocator, nullptr);
+}
+
+TF_Tensor* DoubleTensor(double v) {
+  const int num_bytes = sizeof(double);
+  double* values = new double[1];
+  values[0] = v;
+  return TF_NewTensor(TF_DOUBLE, nullptr, 0, values, num_bytes,
+                      &DoubleDeallocator, nullptr);
+}
+
+TF_Tensor* FloatTensor(float v) {
+  const int num_bytes = sizeof(float);
+  float* values = new float[1];
+  values[0] = v;
+  return TF_NewTensor(TF_FLOAT, nullptr, 0, values, num_bytes,
+                      &FloatDeallocator, nullptr);
+}
+
+// All the *Helper methods are used as a workaround for the restrictions that
+// one cannot call ASSERT_* methods in non-void-returning functions (when
+// exceptions are disabled during compilation)
+void PlaceholderHelper(TF_Graph* graph, TF_Status* s, const char* name,
+                       TF_DataType dtype, const std::vector<int64_t>& dims,
+                       TF_Operation** op) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "Placeholder", name);
+  TF_SetAttrType(desc, "dtype", dtype);
+  if (!dims.empty()) {
+    TF_SetAttrShape(desc, "shape", dims.data(), dims.size());
+  }
+  *op = TF_FinishOperation(desc, s);
+  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  ASSERT_NE(*op, nullptr);
+}
+
+TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, const char* name,
+                          TF_DataType dtype, const std::vector<int64_t>& dims) {
+  TF_Operation* op;
+  PlaceholderHelper(graph, s, name, dtype, dims, &op);
+  return op;
+}
+
+void ConstHelper(TF_Tensor* t, TF_Graph* graph, TF_Status* s, const char* name,
+                 TF_Operation** op) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "Const", name);
+  TF_SetAttrTensor(desc, "value", t, s);
+  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  TF_SetAttrType(desc, "dtype", TF_TensorType(t));
+  *op = TF_FinishOperation(desc, s);
+  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  ASSERT_NE(*op, nullptr);
+}
+
+TF_Operation* Const(TF_Tensor* t, TF_Graph* graph, TF_Status* s,
+                    const char* name) {
+  TF_Operation* op;
+  ConstHelper(t, graph, s, name, &op);
+  return op;
+}
+
+TF_Operation* ScalarConst(bool v, TF_Graph* graph, TF_Status* s,
+                          const char* name) {
+  unique_tensor_ptr tensor(BoolTensor(v), TF_DeleteTensor);
+  return Const(tensor.get(), graph, s, name);
+}
+
+TF_Operation* ScalarConst(int32_t v, TF_Graph* graph, TF_Status* s,
+                          const char* name) {
+  unique_tensor_ptr tensor(Int32Tensor(v), TF_DeleteTensor);
+  return Const(tensor.get(), graph, s, name);
+}
+
+TF_Operation* ScalarConst(double v, TF_Graph* graph, TF_Status* s,
+                          const char* name) {
+  unique_tensor_ptr tensor(DoubleTensor(v), TF_DeleteTensor);
+  return Const(tensor.get(), graph, s, name);
+}
+
+TF_Operation* ScalarConst(float v, TF_Graph* graph, TF_Status* s,
+                          const char* name) {
+  unique_tensor_ptr tensor(FloatTensor(v), TF_DeleteTensor);
+  return Const(tensor.get(), graph, s, name);
+}
+
+void AddOpHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                 TF_Status* s, const char* name, TF_Operation** op,
+                 bool check) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name);
+  TF_Output add_inputs[2] = {{l, 0}, {r, 0}};
+  TF_AddInputList(desc, add_inputs, 2);
+  *op = TF_FinishOperation(desc, s);
+  if (check) {
+    ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+    ASSERT_NE(*op, nullptr);
+  }
+}
+
+TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name) {
+  TF_Operation* op;
+  AddOpHelper(l, r, graph, s, name, &op, true);
+  return op;
+}
+
+TF_Operation* AddNoCheck(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                         TF_Status* s, const char* name) {
+  TF_Operation* op;
+  AddOpHelper(l, r, graph, s, name, &op, false);
+  return op;
+}
+
+TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r,
+                                    TF_Graph* graph, TF_Operation* ctrl_op,
+                                    TF_Status* s, const char* name) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name);
+  TF_Output add_inputs[2] = {{l, 0}, {r, 0}};
+  TF_AddInputList(desc, add_inputs, 2);
+  TF_AddControlInput(desc, ctrl_op);
+  return TF_FinishOperation(desc, s);
+}
+
+// If `op_device` is non-empty, set the created op on that device.
+void BinaryOpHelper(const char* op_name, TF_Operation* l, TF_Operation* r,
+                    TF_Graph* graph, TF_Status* s, const char* name,
+                    TF_Operation** op, const string& op_device, bool check) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, op_name, name);
+  if (!op_device.empty()) {
+    TF_SetDevice(desc, op_device.c_str());
+  }
+  TF_AddInput(desc, {l, 0});
+  TF_AddInput(desc, {r, 0});
+  *op = TF_FinishOperation(desc, s);
+  if (check) {
+    ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+    ASSERT_NE(*op, nullptr);
+  }
+}
+
+TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                            const string& op_device, TF_Status* s,
+                            const char* name) {
+  TF_Operation* op;
+  BinaryOpHelper("Min", l, r, graph, s, name, &op, op_device, true);
+  return op;
+}
+
+TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name) {
+  return MinWithDevice(l, r, graph, /*op_device=*/"", s, name);
+}
+
+TF_Operation* Mul(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name) {
+  TF_Operation* op;
+  BinaryOpHelper("Mul", l, r, graph, s, name, &op, "", true);
+  return op;
+}
+
+TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s,
+                  const char* name) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name);
+  TF_Output inputs[2] = {l, r};
+  TF_AddInputList(desc, inputs, 2);
+  return TF_FinishOperation(desc, s);
+}
+
+void NegHelper(TF_Operation* n, TF_Graph* graph, TF_Status* s, const char* name,
+               TF_Operation** op) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "Neg", name);
+  TF_Output neg_input = {n, 0};
+  TF_AddInput(desc, neg_input);
+  *op = TF_FinishOperation(desc, s);
+  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  ASSERT_NE(*op, nullptr);
+}
+
+TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s,
+                  const char* name) {
+  TF_Operation* op;
+  NegHelper(n, graph, s, name, &op);
+  return op;
+}
+
+TF_Operation* LessThan(TF_Output l, TF_Output r, TF_Graph* graph,
+                       TF_Status* s) {
+  TF_OperationDescription* desc = TF_NewOperation(graph, "Less", "less_than");
+  TF_AddInput(desc, l);
+  TF_AddInput(desc, r);
+  return TF_FinishOperation(desc, s);
+}
+
+TF_Operation* RandomUniform(TF_Operation* shape, TF_DataType dtype,
+                            TF_Graph* graph, TF_Status* s) {
+  TF_OperationDescription* desc =
+      TF_NewOperation(graph, "RandomUniform", "random_uniform");
+  TF_AddInput(desc, {shape, 0});
+  TF_SetAttrType(desc, "dtype", dtype);
+  return TF_FinishOperation(desc, s);
+}
+
+void Split3Helper(TF_Operation* input, TF_Graph* graph, TF_Status* s,
+                  const char* name, TF_Operation** op) {
+  TF_Operation* zero =
+      ScalarConst(0, graph, s, absl::StrCat(name, "_const0").c_str());
+  TF_OperationDescription* desc = TF_NewOperation(graph, "Split", name);
+  TF_AddInput(desc, {zero, 0});
+  TF_AddInput(desc, {input, 0});
+  TF_SetAttrInt(desc, "num_split", 3);
+  TF_SetAttrType(desc, "T", TF_INT32);
+  // Set device to CPU since there is no version of split for int32 on GPU
+  // TODO(iga): Convert all these helpers and tests to use floats because
+  // they are usually available on GPUs. After doing this, remove TF_SetDevice
+  // call in c_api_function_test.cc
+  TF_SetDevice(desc, "/cpu:0");
+  *op = TF_FinishOperation(desc, s);
+  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  ASSERT_NE(*op, nullptr);
+}
+
+TF_Operation* Split3(TF_Operation* input, TF_Graph* graph, TF_Status* s,
+                     const char* name) {
+  TF_Operation* op;
+  Split3Helper(input, graph, s, name, &op);
+  return op;
+}
+
+bool IsPlaceholder(const tensorflow::NodeDef& node_def) {
+  if (node_def.op() != "Placeholder" || node_def.name() != "feed") {
+    return false;
+  }
+  bool found_dtype = false;
+  bool found_shape = false;
+  for (const auto& attr : node_def.attr()) {
+    if (attr.first == "dtype") {
+      if (attr.second.type() == tensorflow::DT_INT32) {
+        found_dtype = true;
+      } else {
+        return false;
+      }
+    } else if (attr.first == "shape") {
+      found_shape = true;
+    }
+  }
+  return found_dtype && found_shape;
+}
+
+bool IsScalarConst(const tensorflow::NodeDef& node_def, int v) {
+  if (node_def.op() != "Const" || node_def.name() != "scalar") {
+    return false;
+  }
+  bool found_dtype = false;
+  bool found_value = false;
+  for (const auto& attr : node_def.attr()) {
+    if (attr.first == "dtype") {
+      if (attr.second.type() == tensorflow::DT_INT32) {
+        found_dtype = true;
+      } else {
+        return false;
+      }
+    } else if (attr.first == "value") {
+      if (attr.second.has_tensor() &&
+          attr.second.tensor().int_val_size() == 1 &&
+          attr.second.tensor().int_val(0) == v) {
+        found_value = true;
+      } else {
+        return false;
+      }
+    }
+  }
+  return found_dtype && found_value;
+}
+
+bool IsAddN(const tensorflow::NodeDef& node_def, int n) {
+  if (node_def.op() != "AddN" || node_def.name() != "add" ||
+      node_def.input_size() != n) {
+    return false;
+  }
+  bool found_t = false;
+  bool found_n = false;
+  for (const auto& attr : node_def.attr()) {
+    if (attr.first == "T") {
+      if (attr.second.type() == tensorflow::DT_INT32) {
+        found_t = true;
+      } else {
+        return false;
+      }
+    } else if (attr.first == "N") {
+      if (attr.second.i() == n) {
+        found_n = true;
+      } else {
+        return false;
+      }
+    }
+  }
+  return found_t && found_n;
+}
+
+bool IsNeg(const tensorflow::NodeDef& node_def, const string& input) {
+  return node_def.op() == "Neg" && node_def.name() == "neg" &&
+         node_def.input_size() == 1 && node_def.input(0) == input;
+}
+
+bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def) {
+  TF_Status* s = TF_NewStatus();
+  TF_Buffer* buffer = TF_NewBuffer();
+  TF_GraphToGraphDef(graph, buffer, s);
+  bool ret = TF_GetCode(s) == TF_OK;
+  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  if (ret) ret = graph_def->ParseFromArray(buffer->data, buffer->length);
+  TF_DeleteBuffer(buffer);
+  TF_DeleteStatus(s);
+  return ret;
+}
+
+bool GetNodeDef(TF_Operation* oper, tensorflow::NodeDef* node_def) {
+  TF_Status* s = TF_NewStatus();
+  TF_Buffer* buffer = TF_NewBuffer();
+  TF_OperationToNodeDef(oper, buffer, s);
+  bool ret = TF_GetCode(s) == TF_OK;
+  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  if (ret) ret = node_def->ParseFromArray(buffer->data, buffer->length);
+  TF_DeleteBuffer(buffer);
+  TF_DeleteStatus(s);
+  return ret;
+}
+
+bool GetFunctionDef(TF_Function* func, tensorflow::FunctionDef* func_def) {
+  TF_Status* s = TF_NewStatus();
+  TF_Buffer* buffer = TF_NewBuffer();
+  TF_FunctionToFunctionDef(func, buffer, s);
+  bool ret = TF_GetCode(s) == TF_OK;
+  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  if (ret) ret = func_def->ParseFromArray(buffer->data, buffer->length);
+  TF_DeleteBuffer(buffer);
+  TF_DeleteStatus(s);
+  return ret;
+}
+
+bool GetAttrValue(TF_Operation* oper, const char* attr_name,
+                  tensorflow::AttrValue* attr_value, TF_Status* s) {
+  TF_Buffer* buffer = TF_NewBuffer();
+  TF_OperationGetAttrValueProto(oper, attr_name, buffer, s);
+  bool ret = TF_GetCode(s) == TF_OK;
+  if (ret) ret = attr_value->ParseFromArray(buffer->data, buffer->length);
+  TF_DeleteBuffer(buffer);
+  return ret;
+}
+
+std::vector<std::pair<string, string>> GetGradDefs(
+    const tensorflow::GraphDef& graph_def) {
+  std::vector<std::pair<string, string>> grads;
+  for (const tensorflow::GradientDef& grad : graph_def.library().gradient()) {
+    grads.emplace_back(grad.function_name(), grad.gradient_func());
+  }
+  std::sort(grads.begin(), grads.end());
+  return grads;
+}
+
+std::vector<string> GetFuncNames(const tensorflow::GraphDef& graph_def) {
+  std::vector<string> names;
+  auto functions = graph_def.library().function();
+  names.reserve(functions.size());
+  for (const tensorflow::FunctionDef& func : functions) {
+    names.push_back(func.signature().name());
+  }
+  std::sort(names.begin(), names.end());
+  return names;
+}
+
+CSession::CSession(TF_Graph* graph, TF_Status* s, bool use_XLA) {
+  TF_SessionOptions* opts = TF_NewSessionOptions();
+  TF_EnableXLACompilation(opts, use_XLA);
+  session_ = TF_NewSession(graph, opts, s);
+  TF_DeleteSessionOptions(opts);
+}
+
+CSession::CSession(TF_Session* session) : session_(session) {}
+
+CSession::~CSession() {
+  TF_Status* s = TF_NewStatus();
+  CloseAndDelete(s);
+  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+  TF_DeleteStatus(s);
+}
+
+void CSession::SetInputs(
+    std::vector<std::pair<TF_Operation*, TF_Tensor*>> inputs) {
+  DeleteInputValues();
+  inputs_.clear();
+  for (const auto& p : inputs) {
+    inputs_.emplace_back(TF_Output{p.first, 0});
+    input_values_.emplace_back(p.second);
+  }
+}
+
+void CSession::SetOutputs(std::initializer_list<TF_Operation*> outputs) {
+  ResetOutputValues();
+  outputs_.clear();
+  for (TF_Operation* o : outputs) {
+    outputs_.emplace_back(TF_Output{o, 0});
+  }
+  output_values_.resize(outputs_.size());
+}
+
+void CSession::SetOutputs(const std::vector<TF_Output>& outputs) {
+  ResetOutputValues();
+  outputs_ = outputs;
+  output_values_.resize(outputs_.size());
+}
+
+void CSession::SetTargets(std::initializer_list<TF_Operation*> targets) {
+  targets_.clear();
+  for (TF_Operation* t : targets) {
+    targets_.emplace_back(t);
+  }
+}
+
+void CSession::Run(TF_Status* s) {
+  if (inputs_.size() != input_values_.size()) {
+    ADD_FAILURE() << "Call SetInputs() before Run()";
+    return;
+  }
+  ResetOutputValues();
+  output_values_.resize(outputs_.size(), nullptr);
+
+  const TF_Output* inputs_ptr = inputs_.empty() ? nullptr : &inputs_[0];
+  TF_Tensor* const* input_values_ptr =
+      input_values_.empty() ? nullptr : &input_values_[0];
+
+  const TF_Output* outputs_ptr = outputs_.empty() ? nullptr : &outputs_[0];
+  TF_Tensor** output_values_ptr =
+      output_values_.empty() ? nullptr : &output_values_[0];
+
+  TF_Operation* const* targets_ptr = targets_.empty() ? nullptr : &targets_[0];
+
+  TF_SessionRun(session_, nullptr, inputs_ptr, input_values_ptr, inputs_.size(),
+                outputs_ptr, output_values_ptr, outputs_.size(), targets_ptr,
+                targets_.size(), nullptr, s);
+
+  DeleteInputValues();
+}
+
+void CSession::CloseAndDelete(TF_Status* s) {
+  DeleteInputValues();
+  ResetOutputValues();
+  if (session_ != nullptr) {
+    TF_CloseSession(session_, s);
+    EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
+    TF_DeleteSession(session_, s);
+    session_ = nullptr;
+  }
+}
+
+void CSession::DeleteInputValues() {
+  for (size_t i = 0; i < input_values_.size(); ++i) {
+    TF_DeleteTensor(input_values_[i]);
+  }
+  input_values_.clear();
+}
+
+void CSession::ResetOutputValues() {
+  for (size_t i = 0; i < output_values_.size(); ++i) {
+    if (output_values_[i] != nullptr) TF_DeleteTensor(output_values_[i]);
+  }
+  output_values_.clear();
+}

tensorflow/c/c_test_util.h

@@ -0,0 +1,165 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_C_TEST_UTIL_H_
+#define TENSORFLOW_C_C_TEST_UTIL_H_
+
+#include "tensorflow/c/c_api.h"
+
+#include <vector>
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+using ::tensorflow::string;
+
+typedef std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)>
+    unique_tensor_ptr;
+
+TF_Tensor* BoolTensor(int32_t v);
+
+// Create a tensor with values of type TF_INT8 provided by `values`.
+TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values);
+
+// Create a tensor with values of type TF_INT32 provided by `values`.
+TF_Tensor* Int32Tensor(const int64_t* dims, int num_dims,
+                       const int32_t* values);
+
+// Create 1 dimensional tensor with values from `values`
+TF_Tensor* Int32Tensor(const std::vector<int32_t>& values);
+
+TF_Tensor* Int32Tensor(int32_t v);
+
+TF_Tensor* DoubleTensor(double v);
+
+TF_Tensor* FloatTensor(float v);
+
+TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s,
+                          const char* name = "feed",
+                          TF_DataType dtype = TF_INT32,
+                          const std::vector<int64_t>& dims = {});
+
+TF_Operation* Const(TF_Tensor* t, TF_Graph* graph, TF_Status* s,
+                    const char* name = "const");
+
+TF_Operation* ScalarConst(bool v, TF_Graph* graph, TF_Status* s,
+                          const char* name = "scalar");
+
+TF_Operation* ScalarConst(int32_t v, TF_Graph* graph, TF_Status* s,
+                          const char* name = "scalar");
+
+TF_Operation* ScalarConst(double v, TF_Graph* graph, TF_Status* s,
+                          const char* name = "scalar");
+
+TF_Operation* ScalarConst(float v, TF_Graph* graph, TF_Status* s,
+                          const char* name = "scalar");
+
+TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name = "add");
+
+TF_Operation* AddNoCheck(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                         TF_Status* s, const char* name = "add");
+
+TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r,
+                                    TF_Graph* graph, TF_Operation* ctrl_op,
+                                    TF_Status* s, const char* name = "add");
+
+TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s,
+                  const char* name = "add");
+
+TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name = "min");
+
+TF_Operation* Mul(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name = "mul");
+
+// If `op_device` is non-empty, set the created op on that device.
+TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                            const string& op_device, TF_Status* s,
+                            const char* name = "min");
+
+TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s,
+                  const char* name = "neg");
+
+TF_Operation* LessThan(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s);
+
+TF_Operation* RandomUniform(TF_Operation* shape, TF_DataType dtype,
+                            TF_Graph* graph, TF_Status* s);
+
+// Split `input` along the first dimension into 3 tensors
+TF_Operation* Split3(TF_Operation* input, TF_Graph* graph, TF_Status* s,
+                     const char* name = "split3");
+
+bool IsPlaceholder(const tensorflow::NodeDef& node_def);
+
+bool IsScalarConst(const tensorflow::NodeDef& node_def, int v);
+
+bool IsAddN(const tensorflow::NodeDef& node_def, int n);
+
+bool IsNeg(const tensorflow::NodeDef& node_def, const string& input);
+
+bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def);
+
+bool GetNodeDef(TF_Operation* oper, tensorflow::NodeDef* node_def);
+
+bool GetFunctionDef(TF_Function* func, tensorflow::FunctionDef* func_def);
+
+bool GetAttrValue(TF_Operation* oper, const char* attr_name,
+                  tensorflow::AttrValue* attr_value, TF_Status* s);
+
+// Returns a sorted vector of std::pair<function_name, gradient_func> from
+// graph_def.library().gradient()
+std::vector<std::pair<string, string>> GetGradDefs(
+    const tensorflow::GraphDef& graph_def);
+
+// Returns a sorted vector of names contained in `grad_def`
+std::vector<string> GetFuncNames(const tensorflow::GraphDef& graph_def);
+
+class CSession {
+ public:
+  CSession(TF_Graph* graph, TF_Status* s, bool use_XLA = false);
+  explicit CSession(TF_Session* session);
+
+  ~CSession();
+
+  void SetInputs(std::vector<std::pair<TF_Operation*, TF_Tensor*>> inputs);
+  void SetOutputs(std::initializer_list<TF_Operation*> outputs);
+  void SetOutputs(const std::vector<TF_Output>& outputs);
+  void SetTargets(std::initializer_list<TF_Operation*> targets);
+
+  void Run(TF_Status* s);
+
+  void CloseAndDelete(TF_Status* s);
+
+  TF_Tensor* output_tensor(int i) { return output_values_[i]; }
+
+  TF_Session* mutable_session() { return session_; }
+
+ private:
+  void DeleteInputValues();
+  void ResetOutputValues();
+
+  TF_Session* session_;
+  std::vector<TF_Output> inputs_;
+  std::vector<TF_Tensor*> input_values_;
+  std::vector<TF_Output> outputs_;
+  std::vector<TF_Tensor*> output_values_;
+  std::vector<TF_Operation*> targets_;
+};
+
+#endif  // TENSORFLOW_C_C_TEST_UTIL_H_

tensorflow/c/checkpoint_reader.cc

@@ -0,0 +1,155 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/checkpoint_reader.h"
+
+#include <utility>
+
+#include "absl/container/flat_hash_set.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/saved_tensor_slice_util.h"
+#include "tensorflow/core/util/tensor_bundle/naming.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+
+namespace tensorflow {
+namespace checkpoint {
+
+class TensorSliceReader;
+
+CheckpointReader::CheckpointReader(const std::string& filename,
+                                   TF_Status* status)
+    : reader_(nullptr),
+      v2_reader_(nullptr),
+      var_to_shape_map_(nullptr),
+      var_to_data_type_map_(nullptr) {
+  // Depending on whether this is a V2 ckpt, initializes "reader_" or
+  // "v2_reader_".
+  std::vector<std::string> v2_path;
+  if (Env::Default()->GetMatchingPaths(MetaFilename(filename), &v2_path).ok() &&
+      !v2_path.empty()) {
+    v2_reader_.reset(
+        new BundleReader(Env::Default(), filename /* prefix to a V2 ckpt */));
+    if (!v2_reader_->status().ok()) {
+      tsl::Set_TF_Status_from_Status(status, v2_reader_->status());
+      return;
+    }
+    auto result = BuildV2VarMaps();
+    var_to_shape_map_.swap(result.first);
+    var_to_data_type_map_.swap(result.second);
+  } else {
+    reader_.reset(new TensorSliceReader(filename));
+    if (!reader_->status().ok()) {
+      tsl::Set_TF_Status_from_Status(status, reader_->status());
+      return;
+    }
+    var_to_shape_map_.reset(
+        new TensorSliceReader::VarToShapeMap(reader_->GetVariableToShapeMap()));
+    var_to_data_type_map_.reset(new TensorSliceReader::VarToDataTypeMap(
+        reader_->GetVariableToDataTypeMap()));
+  }
+}
+
+bool CheckpointReader::HasTensor(const std::string& name) const {
+  if (reader_ != nullptr) {
+    return reader_->HasTensor(name, nullptr, nullptr);
+  }
+  return v2_reader_->Contains(name);
+}
+
+const TensorSliceReader::VarToShapeMap&
+CheckpointReader::GetVariableToShapeMap() const {
+  CHECK(var_to_shape_map_);
+  return *var_to_shape_map_;
+}
+
+const TensorSliceReader::VarToDataTypeMap&
+CheckpointReader::GetVariableToDataTypeMap() const {
+  CHECK(var_to_data_type_map_);
+  return *var_to_data_type_map_;
+}
+
+const std::string CheckpointReader::DebugString() const {
+  if (reader_ != nullptr) return reader_->DebugString();
+  return v2_reader_->DebugString();
+}
+
+void CheckpointReader::GetTensor(
+    const std::string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor,
+    TF_Status* out_status) const {
+  absl::Status status;
+  if (reader_ != nullptr) {
+    status = reader_->GetTensor(name, out_tensor);
+  } else {
+    tensorflow::DataType dtype;
+    tensorflow::TensorShape shape;
+    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
+    if (status.ok()) {
+      out_tensor->reset(new Tensor(dtype, shape));
+      status = v2_reader_->Lookup(name, out_tensor->get());
+      if (!status.ok()) out_tensor->reset();
+    }
+  }
+  if (!status.ok()) {
+    tsl::Set_TF_Status_from_Status(out_status, status);
+  }
+}
+
+std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>,
+          std::unique_ptr<TensorSliceReader::VarToDataTypeMap>>
+CheckpointReader::BuildV2VarMaps() {
+  CHECK(v2_reader_ != nullptr);
+  CHECK(v2_reader_->status().ok());
+
+  // First pass: filters out the entries of the slices.
+  absl::flat_hash_set<std::string> filtered_keys;
+  BundleEntryProto entry;
+  v2_reader_->Seek(kHeaderEntryKey);
+  for (v2_reader_->Next(); v2_reader_->Valid(); v2_reader_->Next()) {
+    CHECK(entry.ParseFromString(v2_reader_->value()))
+        << entry.InitializationErrorString();
+    for (int i = 0; i < entry.slices_size(); ++i) {
+      const auto& slice_proto = entry.slices(i);
+      CHECK(filtered_keys
+                .insert(EncodeTensorNameSlice(
+                    std::string(v2_reader_->key()) /* full var's name */,
+                    TensorSlice(slice_proto)))
+                .second);
+    }
+  }
+
+  // Second pass: adds the entries, ignoring the filtered keys.
+  std::unique_ptr<TensorSliceReader::VarToShapeMap> var_to_shape_map(
+      new TensorSliceReader::VarToShapeMap);
+  std::unique_ptr<TensorSliceReader::VarToDataTypeMap> var_to_data_type_map(
+      new TensorSliceReader::VarToDataTypeMap);
+  v2_reader_->Seek(kHeaderEntryKey);
+  for (v2_reader_->Next(); v2_reader_->Valid(); v2_reader_->Next()) {
+    if (filtered_keys.count(v2_reader_->key()) > 0) continue;
+    CHECK(entry.ParseFromString(v2_reader_->value()))
+        << entry.InitializationErrorString();
+    std::string key(v2_reader_->key());
+    (*var_to_shape_map)[key] = TensorShape(entry.shape());
+    (*var_to_data_type_map)[key] = DataType(entry.dtype());
+  }
+  // The returned pointers are owned by the caller.
+  return std::make_pair(std::move(var_to_shape_map),
+                        std::move(var_to_data_type_map));
+}
+
+}  // namespace checkpoint
+}  // namespace tensorflow

tensorflow/c/checkpoint_reader.h

@@ -0,0 +1,83 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_CHECKPOINT_READER_H_
+#define TENSORFLOW_C_CHECKPOINT_READER_H_
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+#include "tensorflow/core/util/tensor_slice_reader.h"
+
+namespace tensorflow {
+namespace checkpoint {
+
+class TensorSliceReader;
+
+// A wrapper around BundleReader (for V2 checkpoints) and
+// checkpoint::TensorSliceReader (for V1), that is more easily SWIG wrapped for
+// other languages.
+//
+// The class currently only interacts with single-slice (i.e., non-partitioned)
+// variables.
+class CheckpointReader {
+ public:
+  CheckpointReader(const std::string& filename, TF_Status* status);
+
+  bool HasTensor(const std::string& name) const;
+  const std::string DebugString() const;
+
+  // Returns a map from variable names to their shapes.  Slices of a partitioned
+  // tensor are combined into a single entry.
+  const TensorSliceReader::VarToShapeMap& GetVariableToShapeMap() const;
+
+  // Returns a map from variable names to their data types.  Slices of a
+  // partitioned tensor are combined into a single entry.
+  const TensorSliceReader::VarToDataTypeMap& GetVariableToDataTypeMap() const;
+
+  // Attempts to look up the tensor named "name" and stores the found result in
+  // "out_tensor".
+  void GetTensor(const std::string& name,
+                 std::unique_ptr<tensorflow::Tensor>* out_tensor,
+                 TF_Status* out_status) const;
+
+ private:
+  // Uses "v2_reader_" to build "var name -> shape" and "var name -> data type"
+  // maps; both owned by caller.
+  // REQUIRES: "v2_reader_ != nullptr && v2_reader_.status().ok()".
+  std::pair<std::unique_ptr<TensorSliceReader::VarToShapeMap>,
+            std::unique_ptr<TensorSliceReader::VarToDataTypeMap> >
+  BuildV2VarMaps();
+
+  // Invariant: exactly one of "reader_" and "v2_reader_" is non-null.
+  std::unique_ptr<TensorSliceReader> reader_;
+  std::unique_ptr<BundleReader> v2_reader_;
+
+  std::unique_ptr<TensorSliceReader::VarToShapeMap> var_to_shape_map_;
+  std::unique_ptr<TensorSliceReader::VarToDataTypeMap> var_to_data_type_map_;
+
+  CheckpointReader(const CheckpointReader&) = delete;
+  void operator=(const CheckpointReader&) = delete;
+};
+
+}  // namespace checkpoint
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_CHECKPOINT_READER_H_

tensorflow/c/conversion_macros.h

@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_CONVERSION_MACROS_H_
+#define TENSORFLOW_C_CONVERSION_MACROS_H_
+
+#define DEFINE_CONVERSION_FUNCTIONS(cpp_impl, wrapper)                         \
+  inline cpp_impl *unwrap(wrapper *w) {                                        \
+    return reinterpret_cast<cpp_impl *>(w);                                    \
+  }                                                                            \
+                                                                               \
+  inline const cpp_impl *unwrap(const wrapper *w) {                            \
+    return reinterpret_cast<const cpp_impl *>(w);                              \
+  }                                                                            \
+                                                                               \
+  inline wrapper *wrap(cpp_impl *i) { return reinterpret_cast<wrapper *>(i); } \
+  inline const wrapper *wrap(const cpp_impl *i) {                              \
+    return reinterpret_cast<const wrapper *>(i);                               \
+  }
+
+#endif  // TENSORFLOW_C_CONVERSION_MACROS_H_

tensorflow/c/eager/abstract_context.h

@@ -0,0 +1,82 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_
+
+#include <memory>
+
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+
+namespace tensorflow {
+
+// Abstract interface to a context.
+//
+// This serves as a factory for creating `AbstractOperation`s and for
+// registering traced functions.
+// Operations creation within a context can only be executed in that context
+// (for now at least).
+// Implementations of the context may contain some state e.g. an execution
+// environment, a traced representation etc.
+class AbstractContext {
+ protected:
+  enum AbstractContextKind { kGraph, kMlir, kEager, kTfrt, kTape, kOpHandler };
+  explicit AbstractContext(AbstractContextKind kind) : kind_(kind) {}
+  virtual ~AbstractContext() {}
+
+ public:
+  AbstractContextKind getKind() const { return kind_; }
+
+  // Release any underlying resources, including the interface object.
+  //
+  // WARNING: The destructor of this class is marked as protected to disallow
+  // clients from directly destroying this object since it may manage its own
+  // lifetime through ref counting. Thus clients MUST call Release() in order to
+  // destroy an instance of this class.
+  virtual void Release() = 0;
+
+  // Creates an operation builder and ties it to this context.
+  // The returned object can be used for setting operation's attributes,
+  // adding inputs and finally executing (immediately or lazily as in tracing)
+  // it in this context.
+  virtual AbstractOperation* CreateOperation() = 0;
+
+  // Registers a function with this context, after this the function is
+  // available to be called/referenced by its name in this context.
+  virtual absl::Status RegisterFunction(AbstractFunction*) = 0;
+  // Remove a function. 'func' argument is the name of a previously added
+  // FunctionDef. The name is in fdef.signature.name.
+  virtual absl::Status RemoveFunction(const std::string& func) = 0;
+
+ private:
+  const AbstractContextKind kind_;
+};
+
+namespace internal {
+struct AbstractContextDeleter {
+  void operator()(AbstractContext* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+}  // namespace internal
+
+using AbstractContextPtr =
+    std::unique_ptr<AbstractContext, internal::AbstractContextDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_

tensorflow/c/eager/abstract_function.h

@@ -0,0 +1,56 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_
+
+#include "absl/status/statusor.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/intrusive_ptr.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+class FunctionRecord;
+
+// A traced function: this hides the complexity of converting the serialized
+// representation between various supported formats e.g. FunctionDef and Mlir
+// function.
+class AbstractFunction : public core::RefCounted {
+ protected:
+  enum AbstractFunctionKind { kGraph, kMlir };
+  explicit AbstractFunction(AbstractFunctionKind kind) : kind_(kind) {}
+
+ public:
+  // Returns which subclass is this instance of.
+  AbstractFunctionKind getKind() const { return kind_; }
+
+  // Returns the AbstractFunction as a FunctionDef.
+  virtual absl::Status GetFunctionDef(const FunctionDef**) = 0;
+
+  // Returns a shared reference to the wrapped function.
+  virtual absl::StatusOr<core::RefCountPtr<FunctionRecord>>
+  GetFunctionRecord() = 0;
+
+ private:
+  const AbstractFunctionKind kind_;
+};
+
+using AbstractFunctionPtr =
+    tensorflow::core::IntrusivePtr<tensorflow::AbstractFunction>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_

tensorflow/c/eager/abstract_op_attrs.h

@@ -0,0 +1,54 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_OP_ATTRS_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_OP_ATTRS_H_
+
+#include "absl/container/inlined_vector.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Attributes of an op.
+class AbstractOpAttrs {
+ protected:
+  enum AbstractOpAttrsKind { kEager, kTfrt };
+  explicit AbstractOpAttrs(AbstractOpAttrsKind kind) : kind_(kind) {}
+
+ public:
+  // Returns which subclass is this instance of.
+  AbstractOpAttrsKind getKind() const { return kind_; }
+  virtual ~AbstractOpAttrs() = default;
+
+  // Returns the AbstractFunction as a FunctionDef.
+  virtual void GetNameAttrList(
+      tensorflow::NameAttrList* name_and_attrs) const = 0;
+
+  virtual bool GetInt(absl::string_view, int64_t* result) const = 0;
+  virtual bool GetFloat(absl::string_view attr_name, float* result) const = 0;
+  virtual bool GetBool(absl::string_view attr_name, bool* result) const = 0;
+  virtual bool GetType(absl::string_view attr_name, DataType* result) const = 0;
+  virtual absl::Status GetTypeList(
+      absl::string_view attr_name,
+      absl::InlinedVector<DataType, 4>* type_list) const = 0;
+
+ private:
+  const AbstractOpAttrsKind kind_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_OP_ATTRS_H_

tensorflow/c/eager/abstract_operation.h

@@ -0,0 +1,172 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_
+
+#include <memory>
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Abstract interface to an operation.
+// This interface allows building and executing an operation in either
+// tracing or immediate execution mode.
+class AbstractOperation {
+ protected:
+  enum AbstractOperationKind {
+    kGraph,
+    kMlir,
+    kEager,
+    kTfrt,
+    kTape,
+    kOpHandler
+  };
+  explicit AbstractOperation(AbstractOperationKind kind) : kind_(kind) {}
+  virtual ~AbstractOperation() {}
+
+ public:
+  AbstractOperationKind getKind() const { return kind_; }
+
+  // Release any underlying resources, including the interface object.
+  //
+  // WARNING: The destructor of this class is marked as protected to disallow
+  // clients from directly destroying this object since it may manage it's own
+  // lifetime through ref counting. Thus this must be allocated on the heap and
+  // clients MUST call Release() in order to destroy an instance of this class.
+  virtual void Release() = 0;
+
+  virtual absl::Status Reset(const char* op, const char* raw_device_name) = 0;
+
+  virtual const std::string& Name() const = 0;
+
+  // Returns the operation's device name.
+  //
+  // The value returned may be different from the one set by SetDeviceName, but
+  // it will be compatible with it: the name will be updated by device placement
+  // logic to refer to the specific device chosen.
+  //
+  // Example: If one calls `op->SetDeviceName("/device:GPU")`, the value
+  // returned by DeviceName should be "/device:GPU:*" until a particular GPU is
+  // chosen for the operation by the device placement logic in the
+  // executor. After that, the value returned by DeviceName will be a full
+  // device name such as "/job:localhost/replica:0/task:0/device:GPU:1".
+  virtual const std::string& DeviceName() const = 0;
+
+  // Sets the operation device name.
+  //
+  // The given `name` must be parseable by DeviceNameUtils::ParseFullName, and
+  // the result will be used as a constraint for device placement. See the
+  // documentation for DeviceName for more details.
+  //
+  // The value will override the previous value - that is, no "merging" of
+  // existing and given constraints will be performed.
+  virtual absl::Status SetDeviceName(const char* name) = 0;
+
+  virtual absl::Status AddInput(AbstractTensorHandle* input) = 0;
+  virtual absl::Status AddInputList(
+      absl::Span<AbstractTensorHandle* const> inputs) = 0;
+  virtual absl::Status Execute(absl::Span<AbstractTensorHandle*> retvals,
+                               int* num_retvals) = 0;
+
+  virtual absl::Status SetAttrString(const char* attr_name, const char* data,
+                                     size_t length) = 0;
+  virtual absl::Status SetAttrInt(const char* attr_name, int64_t value) = 0;
+  virtual absl::Status SetAttrFloat(const char* attr_name, float value) = 0;
+  virtual absl::Status SetAttrBool(const char* attr_name, bool value) = 0;
+  virtual absl::Status SetAttrType(const char* attr_name, DataType value) = 0;
+  virtual absl::Status SetAttrShape(const char* attr_name, const int64_t* dims,
+                                    const int num_dims) = 0;
+  virtual absl::Status SetAttrShape(const char* attr_name,
+                                    const PartialTensorShape shape);
+  virtual absl::Status SetAttrFunction(const char* attr_name,
+                                       const AbstractOperation* value) = 0;
+  virtual absl::Status SetAttrFunctionName(const char* attr_name,
+                                           const char* value,
+                                           size_t length) = 0;
+  virtual absl::Status SetAttrTensor(const char* attr_name,
+                                     AbstractTensorInterface* tensor) = 0;
+  virtual absl::Status SetAttrStringList(const char* attr_name,
+                                         const void* const* values,
+                                         const size_t* lengths,
+                                         int num_values) = 0;
+  virtual absl::Status SetAttrStringList(const char* attr_name,
+                                         absl::Span<const std::string> values);
+  virtual absl::Status SetAttrFloatList(const char* attr_name,
+                                        const float* values,
+                                        int num_values) = 0;
+  virtual absl::Status SetAttrIntList(const char* attr_name,
+                                      const int64_t* values,
+                                      int num_values) = 0;
+  virtual absl::Status SetAttrTypeList(const char* attr_name,
+                                       const DataType* values,
+                                       int num_values) = 0;
+  virtual absl::Status SetAttrBoolList(const char* attr_name,
+                                       const unsigned char* values,
+                                       int num_values) = 0;
+  virtual absl::Status SetAttrShapeList(const char* attr_name,
+                                        const int64_t** dims,
+                                        const int* num_dims,
+                                        int num_values) = 0;
+  virtual absl::Status SetAttrFunctionList(
+      const char* attr_name, absl::Span<const AbstractOperation*> values) = 0;
+
+ private:
+  const AbstractOperationKind kind_;
+};
+
+// TODO(b/193656009): Defining these in a cc file causes linker errors with
+// fastbuild.
+inline absl::Status AbstractOperation::SetAttrShape(
+    const char* attr_name, const PartialTensorShape shape) {
+  return SetAttrShape(attr_name, shape.dim_sizes().data(), shape.dims());
+}
+
+inline absl::Status AbstractOperation::SetAttrStringList(
+    const char* attr_name, absl::Span<const std::string> values) {
+  std::vector<const char*> raw_strs;
+  std::vector<size_t> lengths;
+  raw_strs.reserve(values.size());
+  lengths.reserve(values.size());
+  for (const auto& s : values) {
+    raw_strs.emplace_back(s.data());
+    lengths.emplace_back(s.size());
+  }
+  return SetAttrStringList(attr_name,
+                           reinterpret_cast<const void**>(raw_strs.data()),
+                           lengths.data(), values.size());
+}
+
+namespace internal {
+struct AbstractOperationDeleter {
+  void operator()(AbstractOperation* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+}  // namespace internal
+
+using AbstractOperationPtr =
+    std::unique_ptr<AbstractOperation, internal::AbstractOperationDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_

tensorflow/c/eager/abstract_tensor_handle.cc

@@ -0,0 +1,40 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+
+namespace tensorflow {
+
+std::string AbstractTensorHandle::DebugString() const {
+  PartialTensorShape shape;
+  absl::Status s = Shape(&shape);
+  std::string shape_string;
+  if (!s.ok()) {
+    shape_string = "<error computing shape>";
+  } else {
+    shape_string = shape.DebugString();
+  }
+  return absl::StrCat("TensorHandle(shape=", shape_string,
+                      ", dtype=", DataType_Name(DataType()),
+                      ", type=", FullType().DebugString(), ")");
+}
+
+absl::Status AbstractTensorHandle::TensorHandleStatus() const {
+  // Tensor handles in current runtime don't carry error info and this method
+  // should always return OK status.
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/c/eager/abstract_tensor_handle.h

@@ -0,0 +1,83 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
+
+#include <memory>
+
+#include "tensorflow/core/framework/full_type.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Abstract interface to a Tensor handle in either tracing or immediate
+// execution mode.
+class AbstractTensorHandle : public core::RefCounted {
+ protected:
+  enum AbstractTensorHandleKind { kGraph, kMlir, kEager, kTfrt, kCustomDevice };
+  explicit AbstractTensorHandle(AbstractTensorHandleKind kind) : kind_(kind) {}
+  ~AbstractTensorHandle() override {}
+
+ public:
+  // Returns tensor dtype.
+  virtual tensorflow::DataType DataType() const = 0;
+
+  // Returns the status of the tensor handle. If it is a tfrt::TensorHandle,
+  // the tensor handle can be an error and return non-OK status.
+  virtual absl::Status TensorHandleStatus() const;
+
+  // Returns tensor shape. If tensor has unknown rank, shape remains untouched.
+  virtual absl::Status Shape(tensorflow::PartialTensorShape* shape) const = 0;
+
+  // Returns tensor (full) type.
+  // While there is no immediate plan to deprecate dtype and shape in favor
+  // of only using full type type information, this is a future possibility.
+  //
+  // Note that map_dtype_to_child_of_tensor() from core/framework/types.h
+  // can be used to set a FullTypeDef based on dtype in a derived class if
+  // appropriate.
+  virtual tensorflow::FullTypeDef FullType() const = 0;
+
+  // The default debug string includes a shape, dtype and FullType.
+  // Implementations are free to override it with something more informative.
+  virtual std::string DebugString() const;
+
+  AbstractTensorHandleKind getKind() const { return kind_; }
+
+ private:
+  const AbstractTensorHandleKind kind_;
+};
+
+namespace internal {
+struct AbstractTensorHandleDeleter {
+  void operator()(AbstractTensorHandle* p) const {
+    if (p != nullptr) {
+      p->Unref();
+    }
+  }
+};
+}  // namespace internal
+
+// TODO(b/185908092): Make AbstractTensorHandlePtr an IntrusivePtr.
+using AbstractTensorHandlePtr =
+    std::unique_ptr<AbstractTensorHandle,
+                    internal::AbstractTensorHandleDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_

tensorflow/c/eager/c_api.h

@@ -0,0 +1,448 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_C_API_H_
+#define TENSORFLOW_C_EAGER_C_API_H_
+
+// C API extensions to experiment with eager execution of kernels.
+// WARNING: Unlike tensorflow/c/c_api.h, the API here is not guaranteed to be
+// stable and can change without notice.
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_macros.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct TFE_ContextOptions TFE_ContextOptions;
+
+// Return a new options object.
+TF_CAPI_EXPORT extern TFE_ContextOptions* TFE_NewContextOptions(void);
+
+// Set the config in TF_ContextOptions.options.
+// config should be a serialized tensorflow.ConfigProto proto.
+// If config was not parsed successfully as a ConfigProto, record the
+// error information in *status.
+TF_CAPI_EXPORT extern void TFE_ContextOptionsSetConfig(
+    TFE_ContextOptions* options, const void* proto, size_t proto_len,
+    TF_Status* status);
+
+// Controls how to act when we try to run an operation on a given device but
+// some input tensors are not on that device.
+// LINT.IfChange
+// Note: Keep in sync with internal copy of enum in eager/context.h.
+typedef enum TFE_ContextDevicePlacementPolicy {
+  // Running operations with input tensors on the wrong device will fail.
+  TFE_DEVICE_PLACEMENT_EXPLICIT = 0,
+  // Copy the tensor to the right device but log a warning.
+  TFE_DEVICE_PLACEMENT_WARN = 1,
+  // Silently copy the tensor, which has a performance cost since the operation
+  // will be blocked till the copy completes. This is the default placement
+  // policy.
+  TFE_DEVICE_PLACEMENT_SILENT = 2,
+  // Placement policy which silently copies int32 tensors but not other dtypes.
+  TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32 = 3,
+} TFE_ContextDevicePlacementPolicy;
+// LINT.ThenChange(//tensorflow/c/eager/immediate_execution_context.h)
+
+// Sets the default execution mode (sync/async). Note that this can be
+// overridden per thread using TFE_ContextSetExecutorForThread.
+TF_CAPI_EXPORT extern void TFE_ContextOptionsSetAsync(TFE_ContextOptions*,
+                                                      unsigned char enable);
+
+TF_CAPI_EXPORT extern void TFE_ContextOptionsSetDevicePlacementPolicy(
+    TFE_ContextOptions*, TFE_ContextDevicePlacementPolicy);
+
+// Destroy an options object.
+TF_CAPI_EXPORT extern void TFE_DeleteContextOptions(TFE_ContextOptions*);
+
+// "Context" under which operations/functions are executed. It encapsulates
+// things like the available devices, resource manager etc.
+// TFE_Context must outlive all tensor handles created using it. In other
+// words, TFE_DeleteContext() must be called after all tensor handles have
+// been deleted (with TFE_DeleteTensorHandle).
+//
+// TODO(ashankar): Merge with TF_Session?
+typedef struct TFE_Context TFE_Context;
+
+TF_CAPI_EXPORT extern TFE_Context* TFE_NewContext(
+    const TFE_ContextOptions* opts, TF_Status* status);
+TF_CAPI_EXPORT extern void TFE_DeleteContext(TFE_Context* ctx);
+TF_CAPI_EXPORT extern TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx,
+                                                            TF_Status* status);
+
+// Clears the internal caches in the TFE context. Useful when reseeding random
+// ops.
+TF_CAPI_EXPORT extern void TFE_ContextClearCaches(TFE_Context* ctx);
+
+// Sets a thread-local device placement policy. After this call, other calls to
+// TFE_Execute in the same thread will use the device policy specified here
+// instead of the device policy used to construct the context. This has no
+// effect on the device policy used by other program threads.
+TF_CAPI_EXPORT extern void TFE_ContextSetThreadLocalDevicePlacementPolicy(
+    TFE_Context* ctx, TFE_ContextDevicePlacementPolicy policy);
+
+// Returns the device placement policy to be used by this context in the current
+// thread.
+TF_CAPI_EXPORT extern TFE_ContextDevicePlacementPolicy
+TFE_ContextGetDevicePlacementPolicy(TFE_Context* ctx);
+
+// A tensorflow.ServerDef specifies remote workers (in addition to the current
+// workers name). Operations created in this context can then be executed on
+// any of these remote workers by setting an appropriate device.
+//
+// If the following is set, all servers identified by the
+// ServerDef must be up when the context is created.
+TF_CAPI_EXPORT extern void TFE_ContextSetServerDef(TFE_Context* ctx,
+                                                   int keep_alive_secs,
+                                                   const void* proto,
+                                                   size_t proto_len,
+                                                   TF_Status* status);
+
+// A handle to a tensor on a device.
+//
+// Like a TF_Tensor, a TFE_TensorHandle refers to a tensor with a value, shape,
+// type etc. Unlike a TF_Tensor, a TFE_TensorHandle may refer to such tensors
+// placed in the memory of different devices or remote address spaces.
+typedef struct TFE_TensorHandle TFE_TensorHandle;
+
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandle(const TF_Tensor* t,
+                                                            TF_Status* status);
+// Indicates that the caller will not be using `h` any more.
+TF_CAPI_EXPORT extern void TFE_DeleteTensorHandle(TFE_TensorHandle* h);
+TF_CAPI_EXPORT extern TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h);
+// This function will block till the operation that produces `h` has completed.
+TF_CAPI_EXPORT extern int TFE_TensorHandleNumDims(TFE_TensorHandle* h,
+                                                  TF_Status* status);
+TF_CAPI_EXPORT extern int64_t TFE_TensorHandleNumElements(TFE_TensorHandle* h,
+                                                          TF_Status* status);
+// This function will block till the operation that produces `h` has completed.
+TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h,
+                                                  int dim_index,
+                                                  TF_Status* status);
+
+// Returns the device of the operation that produced `h`. If `h` was produced by
+// a copy, returns the destination device of the copy. Note that the returned
+// device name is not always the device holding the tensor handle's memory. If
+// you want the latter, use TFE_TensorHandleBackingDeviceName. This function
+// will block till the operation that produces `h` has completed.
+TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName(
+    TFE_TensorHandle* h, TF_Status* status);
+
+// Returns the name of the device in whose memory `h` resides.
+//
+// This function will block till the operation that produces `h` has completed.
+TF_CAPI_EXPORT extern const char* TFE_TensorHandleBackingDeviceName(
+    TFE_TensorHandle* h, TF_Status* status);
+
+// Return a pointer to a new TFE_TensorHandle that shares the underlying tensor
+// with `h`. On success, `status` is set to OK. On failure, `status` reflects
+// the error and a nullptr is returned.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopySharingTensor(
+    TFE_TensorHandle* h, TF_Status* status);
+
+// This function will block till the operation that produces `h` has
+// completed. The memory returned might alias the internal memory used by
+// TensorFlow. Hence, callers should not mutate this memory (for example by
+// modifying the memory region pointed to by TF_TensorData() on the returned
+// TF_Tensor).
+TF_CAPI_EXPORT extern TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h,
+                                                         TF_Status* status);
+
+// Create a new TFE_TensorHandle with the same contents as 'h' but placed
+// in the memory of the device name 'device_name'.
+// If source and destination are the same device, then this creates a new handle
+// that shares the underlying buffer. Otherwise, it currently requires at least
+// one of the source or destination devices to be CPU (i.e., for the source or
+// destination tensor to be placed in host memory).
+// If async execution is enabled, the copy may be enqueued and the call will
+// return "non-ready" handle. Else, this function returns after the copy has
+// been done.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopyToDevice(
+    TFE_TensorHandle* h, TFE_Context* ctx, const char* device_name,
+    TF_Status* status);
+
+// Debugging/Profiling information for TFE_TensorHandle
+//
+// TFE_TensorDebugInfo contains information useful for debugging and
+// profiling tensors.
+typedef struct TFE_TensorDebugInfo TFE_TensorDebugInfo;
+
+// Retrieves TFE_TensorDebugInfo for `handle`.
+// If TFE_TensorHandleTensorDebugInfo succeeds, `status` is set to OK and caller
+// is responsible for deleting returned TFE_TensorDebugInfo.
+// If TFE_TensorHandleTensorDebugInfo fails, `status` is set to appropriate
+// error and nullptr is returned. This function can block till the operation
+// that produces `handle` has completed.
+TF_CAPI_EXPORT extern TFE_TensorDebugInfo* TFE_TensorHandleTensorDebugInfo(
+    TFE_TensorHandle* h, TF_Status* status);
+
+// Deletes `debug_info`.
+TF_CAPI_EXPORT extern void TFE_DeleteTensorDebugInfo(
+    TFE_TensorDebugInfo* debug_info);
+
+// Returns the number of dimensions used to represent the tensor on its device.
+// The number of dimensions used to represent the tensor on device can be
+// different from the number returned by TFE_TensorHandleNumDims.
+// The return value was current at the time of TFE_TensorDebugInfo creation.
+TF_CAPI_EXPORT extern int TFE_TensorDebugInfoOnDeviceNumDims(
+    TFE_TensorDebugInfo* debug_info);
+
+// Returns the number of elements in dimension `dim_index`.
+// Tensor representation on device can be transposed from its representation
+// on host. The data contained in dimension `dim_index` on device
+// can correspond to the data contained in another dimension in on-host
+// representation. The dimensions are indexed using the standard TensorFlow
+// major-to-minor order (slowest varying dimension first),
+// not the XLA's minor-to-major order.
+// On-device dimensions can be padded. TFE_TensorDebugInfoOnDeviceDim returns
+// the number of elements in a dimension after padding.
+// The return value was current at the time of TFE_TensorDebugInfo creation.
+TF_CAPI_EXPORT extern int64_t TFE_TensorDebugInfoOnDeviceDim(
+    TFE_TensorDebugInfo* debug_info, int dim_index);
+
+// Description of the TensorFlow op to execute.
+//
+// Assumes that the provided 'ctx' outlives the returned TFE_Op, i.e.,
+// TFE_DeleteOp() is called before TFE_DeleteContext().
+//
+// Very similar to TF_OperationDescription with some differences:
+// (1) TF_Output or TFE_TensorHandle* as arguments to TF_AddInput,
+//     TF_AddInputList
+// (2) TF_ColocateWith, TF_AddControlInput etc. do not make sense.
+// (3) Implementation detail: Avoid use of NodeBuilder/NodeDefBuilder since
+//     the additional sanity checks there seem unnecessary;
+typedef struct TFE_Op TFE_Op;
+
+TF_CAPI_EXPORT extern TFE_Op* TFE_NewOp(TFE_Context* ctx,
+                                        const char* op_or_function_name,
+                                        TF_Status* status);
+TF_CAPI_EXPORT extern void TFE_DeleteOp(TFE_Op* op);
+
+// Returns the op or function name `op` will execute.
+//
+// The returned string remains valid throughout the lifetime of 'op'.
+TF_CAPI_EXPORT extern const char* TFE_OpGetName(const TFE_Op* op,
+                                                TF_Status* status);
+TF_CAPI_EXPORT extern TFE_Context* TFE_OpGetContext(const TFE_Op* op,
+                                                    TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_OpSetDevice(TFE_Op* op, const char* device_name,
+                                           TF_Status* status);
+// The returned string remains valid throughout the lifetime of 'op'.
+TF_CAPI_EXPORT extern const char* TFE_OpGetDevice(const TFE_Op* op,
+                                                  TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* input,
+                                          TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_OpAddInputList(TFE_Op* op,
+                                              TFE_TensorHandle** inputs,
+                                              int num_inputs,
+                                              TF_Status* status);
+
+// Fetches the current number of inputs attached to `op`.
+//
+// Does not use the operation's definition to determine how many inputs should
+// be attached. It is intended for use with TFE_OpGetFlatInput to inspect an
+// already-finalized operation.
+//
+// Note that TFE_OpGetFlatInputCount and TFE_OpGetFlatInput operate on a flat
+// sequence of inputs, unlike TFE_OpGetInputLength (for getting the length of a
+// particular named input list, which may only be part of the op's inputs).
+TF_CAPI_EXPORT extern int TFE_OpGetFlatInputCount(const TFE_Op* op,
+                                                  TF_Status* status);
+// Returns a borrowed reference to one of `op`'s inputs. Use
+// `TFE_TensorHandleCopySharingTensor` to make a new reference.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_OpGetFlatInput(const TFE_Op* op,
+                                                           int index,
+                                                           TF_Status* status);
+
+TF_CAPI_EXPORT extern TF_AttrType TFE_OpGetAttrType(TFE_Op* op,
+                                                    const char* attr_name,
+                                                    unsigned char* is_list,
+                                                    TF_Status* status);
+// Get an attribute type given an op name; a fusion of TFE_NewOp and
+// TFE_OpGetAttrType for use from Python without the overhead of the individual
+// calls and memory management of TFE_Op.
+TF_CAPI_EXPORT extern TF_AttrType TFE_OpNameGetAttrType(
+    TFE_Context* ctx, const char* op_or_function_name, const char* attr_name,
+    unsigned char* is_list, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_OpSetAttrString(TFE_Op* op,
+                                               const char* attr_name,
+                                               const void* value,
+                                               size_t length);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name,
+                                            int64_t value);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name,
+                                              float value);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrBool(TFE_Op* op, const char* attr_name,
+                                             unsigned char value);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name,
+                                             TF_DataType value);
+// If the number of dimensions is unknown, `num_dims` must be set to
+// -1 and `dims` can be null.  If a dimension is unknown, the
+// corresponding entry in the `dims` array must be -1.
+TF_CAPI_EXPORT extern void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name,
+                                              const int64_t* dims,
+                                              const int num_dims,
+                                              TF_Status* out_status);
+
+// Sets the attribute attr_name to be a function specified by 'function'.
+//
+// TODO(ashankar,iga): Add this functionality to the C API for graph
+// construction. Perhaps we want an AttrValueMap equivalent in the C API?
+TF_CAPI_EXPORT extern void TFE_OpSetAttrFunction(TFE_Op* op,
+                                                 const char* attr_name,
+                                                 const TFE_Op* value);
+
+TF_CAPI_EXPORT void TFE_OpSetAttrFunctionName(TFE_Op* op, const char* attr_name,
+                                              const char* data, size_t length);
+
+TF_CAPI_EXPORT extern void TFE_OpSetAttrTensor(TFE_Op* op,
+                                               const char* attr_name,
+                                               TF_Tensor* tensor,
+                                               TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_OpSetAttrStringList(TFE_Op* op,
+                                                   const char* attr_name,
+                                                   const void* const* values,
+                                                   const size_t* lengths,
+                                                   int num_values);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrIntList(TFE_Op* op,
+                                                const char* attr_name,
+                                                const int64_t* values,
+                                                int num_values);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrFloatList(TFE_Op* op,
+                                                  const char* attr_name,
+                                                  const float* values,
+                                                  int num_values);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrBoolList(TFE_Op* op,
+                                                 const char* attr_name,
+                                                 const unsigned char* values,
+                                                 int num_values);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrTypeList(TFE_Op* op,
+                                                 const char* attr_name,
+                                                 const TF_DataType* values,
+                                                 int num_values);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrShapeList(
+    TFE_Op* op, const char* attr_name, const int64_t** dims,
+    const int* num_dims, int num_values, TF_Status* out_status);
+TF_CAPI_EXPORT extern void TFE_OpSetAttrFunctionList(TFE_Op* op,
+                                                     const char* attr_name,
+                                                     const TFE_Op** value,
+                                                     int num_values);
+
+// Returns the length (number of tensors) of the input argument `input_name`
+// found in the provided `op`.
+TF_CAPI_EXPORT extern int TFE_OpGetInputLength(TFE_Op* op,
+                                               const char* input_name,
+                                               TF_Status* status);
+
+// Returns the length (number of tensors) of the output argument `output_name`
+// found in the provided `op`.
+TF_CAPI_EXPORT extern int TFE_OpGetOutputLength(TFE_Op* op,
+                                                const char* output_name,
+                                                TF_Status* status);
+
+// Execute the operation defined by 'op' and return handles to computed
+// tensors in `retvals`.
+//
+// 'retvals' must point to a pre-allocated array of TFE_TensorHandle* and
+// '*num_retvals' should be set to the size of this array. It is an error if
+// the size of 'retvals' is less than the number of outputs. This call sets
+// *num_retvals to the number of outputs.
+//
+// If async execution is enabled, the call may simply enqueue the execution
+// and return "non-ready" handles in `retvals`. Note that any handles contained
+// in 'op' should not be mutated till the kernel execution actually finishes.
+//
+// For sync execution, if any of the inputs to `op` are not ready, this call
+// will block till they become ready and then return when the kernel execution
+// is done.
+// TODO(agarwal): change num_retvals to int from int*.
+TF_CAPI_EXPORT extern void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals,
+                                       int* num_retvals, TF_Status* status);
+
+// Add a function (serialized FunctionDef protocol buffer) to ctx so
+// that it can be invoked using TFE_Execute.
+TF_CAPI_EXPORT extern void TFE_ContextAddFunctionDef(
+    TFE_Context* ctx, const char* serialized_function_def, size_t size,
+    TF_Status* status);
+
+// Adds a function (created from TF_GraphToFunction or
+// TF_FunctionImportFunctionDef) to the context, allowing it to be executed with
+// TFE_Execute by creating an op with the same name as the function.
+TF_CAPI_EXPORT extern void TFE_ContextAddFunction(TFE_Context* ctx,
+                                                  TF_Function* function,
+                                                  TF_Status* status);
+
+// Removes a function from the context. Once removed, you can no longer
+// TFE_Execute it or TFE_Execute any TFE_Op which has it as an attribute or any
+// other function which calls it as an attribute.
+TF_CAPI_EXPORT extern void TFE_ContextRemoveFunction(TFE_Context* ctx,
+                                                     const char* name,
+                                                     TF_Status* status);
+
+// Checks whether a function is registered under `name`.
+TF_CAPI_EXPORT unsigned char TFE_ContextHasFunction(TFE_Context* ctx,
+                                                    const char* name);
+
+// Enables tracing of RunMetadata on the ops executed from this context.
+TF_CAPI_EXPORT extern void TFE_ContextEnableRunMetadata(TFE_Context* ctx);
+
+// Disables tracing of RunMetadata on the ops executed from this context.
+TF_CAPI_EXPORT extern void TFE_ContextDisableRunMetadata(TFE_Context* ctx);
+
+// Populates the passed-in buffer with a serialized RunMetadata protocol buffer
+// containing any run metadata information accumulated so far and clears this
+// information.
+// If async mode is enabled, this call blocks till all currently pending ops are
+// done.
+TF_CAPI_EXPORT extern void TFE_ContextExportRunMetadata(TFE_Context* ctx,
+                                                        TF_Buffer* buf,
+                                                        TF_Status* status);
+
+// Some TF ops need a step container to be set to limit the lifetime of some
+// resources (mostly TensorArray and Stack, used in while loop gradients in
+// graph mode). Calling this on a context tells it to start a step.
+TF_CAPI_EXPORT extern void TFE_ContextStartStep(TFE_Context* ctx);
+
+// Ends a step. When there is no active step (that is, every started step has
+// been ended) step containers will be cleared. Note: it is not safe to call
+// TFE_ContextEndStep while ops that rely on the step container may be running.
+TF_CAPI_EXPORT extern void TFE_ContextEndStep(TFE_Context* ctx);
+
+#ifdef __cplusplus
+} /* end extern "C" */
+#endif
+
+#ifdef __cplusplus
+// A workaround to ease conversion to and from numpy objects and
+// TFE_TensorHandle's.
+//
+// TODO(ashankar): Figure out an alternative scheme that precludes the need for
+// these API-boundary breaking methods.
+namespace tensorflow {
+class Tensor;
+}  // namespace tensorflow
+
+TFE_TensorHandle* TFE_NewTensorHandle(const tensorflow::Tensor& t,
+                                      TF_Status* status);
+#endif
+
+#endif  // TENSORFLOW_C_EAGER_C_API_H_

tensorflow/c/eager/c_api_cluster_test.cc

@@ -0,0 +1,477 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+namespace {
+
+using ::tensorflow::string;
+
+void ReplaceTaskInServerDef(tensorflow::ServerDef* server_def, int task_index) {
+  tensorflow::JobDef* job_def = server_def->mutable_cluster()->mutable_job(0);
+  int port = tensorflow::testing::PickUnusedPortOrDie();
+  job_def->mutable_tasks()->at(task_index) = absl::StrCat("localhost:", port);
+}
+
+void CheckTFE_TensorHandleHasFloats(TFE_TensorHandle* handle,
+                                    const std::vector<float>& expected_values) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TF_Tensor* t = TFE_TensorHandleResolve(handle, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  std::unique_ptr<float[]> actual_values(new float[expected_values.size()]);
+  EXPECT_EQ(sizeof(float) * expected_values.size(), TF_TensorByteSize(t));
+  memcpy(actual_values.get(), TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+
+  for (int i = 0; i < expected_values.size(); i++) {
+    EXPECT_EQ(expected_values[i], actual_values[i])
+        << "Mismatch in expected values at (zero-based) index " << i;
+  }
+}
+
+void CheckRemoteMatMulExecutesOK(TFE_Context* ctx,
+                                 const char* remote_device_name,
+                                 const char* local_device_name) {
+  TF_Status* status = TF_NewStatus();
+  TFE_TensorHandle* h0_task0 = TestMatrixTensorHandle(ctx);
+
+  TFE_Op* matmul = MatMulOp(ctx, h0_task0, h0_task0);
+  TFE_OpSetDevice(matmul, remote_device_name, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* retvals[1];
+  int num_retvals = 1;
+  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  auto* retval_task0 =
+      TFE_TensorHandleCopyToDevice(retvals[0], ctx, local_device_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  CheckTFE_TensorHandleHasFloats(retval_task0, {7, 10, 15, 22});
+
+  TFE_DeleteTensorHandle(retval_task0);
+  TFE_DeleteTensorHandle(h0_task0);
+  TFE_DeleteTensorHandle(retvals[0]);
+
+  TFE_DeleteOp(matmul);
+
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+  TF_DeleteStatus(status);
+}
+
+// Read the value of variable `var` and save it into `out_value`.
+void ReadVariable(TFE_Context* ctx, TFE_TensorHandle* var,
+                  TFE_TensorHandle** out_value) {
+  TF_Status* status = TF_NewStatus();
+  TFE_Op* op = TFE_NewOp(ctx, "ReadVariableOp", status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  TFE_OpAddInput(op, var, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  int num_retvals = 1;
+  TFE_Execute(op, out_value, &num_retvals, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteOp(op);
+  TF_DeleteStatus(status);
+}
+
+void TestRemoteExecuteChangeServerDef(bool async) {
+  tensorflow::ServerDef server_def = GetServerDef(2);
+
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  const char remote_device_name[] =
+      "/job:localhost/replica:0/task:1/device:CPU:0";
+  const char local_device_name[] =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);
+
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+
+  // Update the server def with a new set of names (worker instead of
+  // localhost).
+  tensorflow::ServerDef updated_server_def = GetServerDef("worker", 2);
+  serialized = updated_server_def.SerializeAsString();
+
+  updated_server_def.set_task_index(1);
+  absl::Status s = tensorflow::GrpcServer::Create(
+      updated_server_def, tensorflow::Env::Default(), &worker_server);
+  ASSERT_TRUE(s.ok()) << s.message();
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // Create a new tensor_handle.
+  TFE_TensorHandle* h0_task0_new = TestMatrixTensorHandle(ctx);
+
+  // Check that copying it to the old remote device (named localhost) fails.
+  TFE_TensorHandleCopyToDevice(h0_task0_new, ctx, remote_device_name, status);
+  EXPECT_NE(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // Copying and executing on the new remote device works.
+  const char new_remote_device_name[] =
+      "/job:worker/replica:0/task:1/device:CPU:0";
+  const char new_local_device_name[] =
+      "/job:worker/replica:0/task:0/device:CPU:0";
+
+  auto* h0_task1_new = TFE_TensorHandleCopyToDevice(
+      h0_task0_new, ctx, new_remote_device_name, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_DeleteTensorHandle(h0_task0_new);
+  TFE_DeleteTensorHandle(h0_task1_new);
+
+  CheckRemoteMatMulExecutesOK(ctx, new_remote_device_name,
+                              new_local_device_name);
+
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+
+  TF_DeleteStatus(status);
+
+  TFE_DeleteContext(ctx);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+}
+
+TEST(CAPI, RemoteExecuteChangeServerDef) {
+  TestRemoteExecuteChangeServerDef(false);
+}
+TEST(CAPI, RemoteExecuteChangeServerDefAsync) {
+  TestRemoteExecuteChangeServerDef(true);
+}
+
+void TestRemoteExecuteUpdateServerDef(bool async) {
+  tensorflow::ServerDef server_def = GetServerDef(2);
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  const char local_device_name[] =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  const char remote_device_name[] =
+      "/job:localhost/replica:0/task:1/device:CPU:0";
+  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);
+
+  TFE_ContextUpdateServerDef(ctx, 0, serialized.data(), serialized.size(),
+                             status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);
+
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+}
+
+TEST(CAPI, RemoteExecuteUpdateServerDef) {
+  TestRemoteExecuteUpdateServerDef(false);
+}
+
+TEST(CAPI, RemoteExecuteUpdateServerDefAsync) {
+  TestRemoteExecuteUpdateServerDef(true);
+}
+
+void TestRemoteExecuteUpdateServerDefResourceAccess(bool async) {
+  tensorflow::ServerDef server_def = GetServerDef(2);
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  const char dev0_name[] = "/job:localhost/replica:0/task:0/device:CPU:0";
+  const char dev1_name[] = "/job:localhost/replica:0/task:1/device:CPU:0";
+
+  TFE_TensorHandle* var_handle0 = TestVariable(ctx, 1.0, dev0_name);
+  EXPECT_NE(var_handle0, nullptr);
+  TFE_TensorHandle* var_handle1 = TestVariable(ctx, 2.0, dev1_name);
+  EXPECT_NE(var_handle1, nullptr);
+
+  TFE_TensorHandle* value_handle = nullptr;
+  ReadVariable(ctx, var_handle1, &value_handle);
+  CheckTFE_TensorHandleHasFloats(value_handle, {2});
+  TFE_DeleteTensorHandle(value_handle);
+
+  // Start a new worker to replace task:1
+  ReplaceTaskInServerDef(&server_def, 1);
+  server_def.set_task_index(1);
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  // Update server def to replace the remote device with the device info on the
+  // new worker (different incarnation ID).
+  server_def.set_task_index(0);
+  string serialized_update = server_def.SerializeAsString();
+  TFE_ContextUpdateServerDef(ctx, 0, serialized_update.data(),
+                             serialized_update.size(), status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // The device of var_handle0 is local device which is the same before and
+  // after cluster update. Remove resource with valid device should succeed.
+  TFE_Op* op = TFE_NewOp(ctx, "DestroyResourceOp", status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, var_handle0, status);
+  TFE_OpSetDevice(op, dev0_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  int num_retvals = 0;
+  TFE_Execute(op, nullptr, &num_retvals, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteOp(op);
+
+  // The device of var_handle1 is remote device, which was replaced during
+  // cluster update. Removing resource with invalid device should fail
+  // gracefully (i.e., with error status) instead of crashing with segfaults.
+  op = TFE_NewOp(ctx, "DestroyResourceOp", status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, var_handle1, status);
+  TFE_OpSetDevice(op, dev1_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  num_retvals = 0;
+  TFE_Execute(op, nullptr, &num_retvals, status);
+  EXPECT_NE(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteOp(op);
+
+  TFE_DeleteTensorHandle(var_handle0);
+  TFE_DeleteTensorHandle(var_handle1);
+
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+}
+
+TEST(CAPI, TestRemoteExecuteUpdateServerDefResourceAccess) {
+  TestRemoteExecuteUpdateServerDefResourceAccess(false);
+}
+
+TEST(CAPI, TestRemoteExecuteUpdateServerDefResourceAccessAsync) {
+  TestRemoteExecuteUpdateServerDefResourceAccess(true);
+}
+
+void TestRemoteExecuteUpdateServerDefWithFailures(bool async) {
+  // Fail fast on GetStatus requests so we can get errors instead of timeout
+  // when updating cluster with non-exsitent worker
+  tensorflow::setenv("GRPC_FAIL_FAST", "TRUE", /*overwrite=*/1);
+
+  tensorflow::ServerDef server_def = GetServerDef(2);
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  const char local_device_name[] =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  const char remote_device_name[] =
+      "/job:localhost/replica:0/task:1/device:CPU:0";
+  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);
+
+  // Adding a non-existent remote worker to cluster def. This should cause the
+  // UpdateServerDef call to fail.
+  tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->mutable_job(0);
+  int port = tensorflow::testing::PickUnusedPortOrDie();
+  job_def->mutable_tasks()->insert({2, absl::StrCat("localhost:", port)});
+  server_def.set_task_index(0);
+  string serialized_update = server_def.SerializeAsString();
+  TFE_ContextUpdateServerDef(ctx, 0, serialized_update.data(),
+                             serialized_update.size(), status);
+  EXPECT_NE(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // Even after the prevoiusly failed cluster update, another update and op
+  // execution should work fine as long as the provided server_def is valid.
+  TFE_ContextUpdateServerDef(ctx, 0, serialized.data(), serialized.size(),
+                             status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);
+
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+  tensorflow::unsetenv("GRPC_FAIL_FAST");
+}
+
+TEST(CAPI, RemoteExecuteUpdateServerDefWithFailures) {
+  TestRemoteExecuteUpdateServerDefWithFailures(false);
+}
+
+TEST(CAPI, RemoteExecuteUpdateServerDefWithFailuresAsync) {
+  TestRemoteExecuteUpdateServerDefWithFailures(true);
+}
+
+void TestConnectToCluster(bool keep_localhost_for_first_connect) {
+  // Fail fast on GetStatus requests so we can get errors instead of timeout
+  // when updating cluster with non-exsitent worker
+  tensorflow::setenv("GRPC_FAIL_FAST", "TRUE", /*overwrite=*/1);
+
+  const string first_name =
+      keep_localhost_for_first_connect ? "localhost" : "abc";
+  tensorflow::ServerDef server_def = GetServerDef(first_name, 1);
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  const string dev0_name = "/job:localhost/replica:0/task:0/device:CPU:0";
+  TFE_TensorHandle* var_handle0 = TestVariable(ctx, 1.0, dev0_name);
+  EXPECT_NE(var_handle0, nullptr);
+
+  absl::Status status2;
+  EXPECT_EQ(tensorflow::unwrap(var_handle0)->DeviceName(&status2), dev0_name);
+
+  // Rename local device
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  const string dev1_name =
+      absl::StrCat("/job:", first_name, "/replica:0/task:0/device:CPU:0");
+  TFE_TensorHandle* var_handle1 = TestVariable(ctx, 2.0, dev1_name);
+  EXPECT_NE(var_handle1, nullptr);
+  EXPECT_EQ(tensorflow::unwrap(var_handle1)->DeviceName(&status2), dev1_name);
+
+  // Another renaming of local device
+  const string second_name = "def";
+  server_def.set_job_name(second_name);
+  server_def.mutable_cluster()->mutable_job(0)->set_name(second_name);
+  (*server_def.mutable_cluster()->mutable_job(0)->mutable_tasks())[0] =
+      absl::StrCat(second_name, ":",
+                   tensorflow::testing::PickUnusedPortOrDie());
+
+  serialized = server_def.SerializeAsString();
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  const string dev2_name = "/job:def/replica:0/task:0/device:CPU:0";
+  TFE_TensorHandle* var_handle2 = TestVariable(ctx, 2.0, dev2_name);
+  EXPECT_NE(var_handle2, nullptr);
+  EXPECT_EQ(tensorflow::unwrap(var_handle2)->DeviceName(&status2), dev2_name);
+
+  TFE_DeleteTensorHandle(var_handle0);
+  TFE_DeleteTensorHandle(var_handle1);
+  TFE_DeleteTensorHandle(var_handle2);
+
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  tensorflow::unsetenv("GRPC_FAIL_FAST");
+}
+
+TEST(CAPI, ConnectToClusterLocalhostFirst) { TestConnectToCluster(false); }
+
+TEST(CAPI, ConnectToClusterRenameFirst) { TestConnectToCluster(true); }
+
+}  // namespace

tensorflow/c/eager/c_api_debug.cc

@@ -0,0 +1,86 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/tfe_tensor_debug_info_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/platform/status.h"
+
+using tensorflow::string;
+
+namespace {
+
+std::vector<int64_t> TensorShapeAsVector(const tensorflow::TensorHandle& handle,
+                                         absl::Status* status) {
+  std::vector<int64_t> shape;
+  int rank = -1;
+  *status = handle.NumDims(&rank);
+  if (!status->ok()) {
+    return shape;
+  }
+  shape.reserve(rank);
+  for (int i = 0; i < rank; ++i) {
+    int64_t dim;
+    *status = handle.Dim(i, &dim);
+    if (!status->ok()) {
+      return shape;
+    }
+    shape.push_back(dim);
+  }
+  return shape;
+}
+
+}  // namespace
+
+extern "C" {
+
+TF_CAPI_EXPORT extern TFE_TensorDebugInfo* TFE_TensorHandleTensorDebugInfo(
+    TFE_TensorHandle* h, TF_Status* status) {
+  tensorflow::TensorHandle* handle =
+      TensorHandleFromInterface(tensorflow::unwrap(h));
+  const tensorflow::Tensor* tensor;
+  status->status = handle->Tensor(&tensor);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  std::vector<int64_t> dev_dims = TensorShapeAsVector(*handle, &status->status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+  return new TFE_TensorDebugInfo(dev_dims);
+}
+
+TF_CAPI_EXPORT extern void TFE_DeleteTensorDebugInfo(
+    TFE_TensorDebugInfo* debug_info) {
+  delete debug_info;
+}
+
+TF_CAPI_EXPORT extern int TFE_TensorDebugInfoOnDeviceNumDims(
+    TFE_TensorDebugInfo* debug_info) {
+  return debug_info->dev_dims.size();
+}
+
+TF_CAPI_EXPORT extern int64_t TFE_TensorDebugInfoOnDeviceDim(
+    TFE_TensorDebugInfo* debug_info, int dim_index) {
+  return debug_info->dev_dims[dim_index];
+}
+
+}  // extern "C"

tensorflow/c/eager/c_api_debug_test.cc

@@ -0,0 +1,62 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api.h"
+
+#include <string.h>
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+
+TEST(CApiDebug, ScalarCPU) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_TensorHandle* h = TestScalarTensorHandle(ctx, 1.0f);
+  TFE_TensorDebugInfo* debug_info = TFE_TensorHandleTensorDebugInfo(h, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  ASSERT_EQ(0, TFE_TensorDebugInfoOnDeviceNumDims(debug_info));
+
+  TFE_DeleteTensorDebugInfo(debug_info);
+  TFE_DeleteTensorHandle(h);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+}
+
+TEST(CApiDebug, 2DCPU) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_TensorHandle* h = TestMatrixTensorHandle3X2(ctx);
+  TFE_TensorDebugInfo* debug_info = TFE_TensorHandleTensorDebugInfo(h, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  ASSERT_EQ(2, TFE_TensorDebugInfoOnDeviceNumDims(debug_info));
+  // Shape is the same for CPU tensors.
+  EXPECT_EQ(3, TFE_TensorDebugInfoOnDeviceDim(debug_info, 0));
+  EXPECT_EQ(2, TFE_TensorDebugInfoOnDeviceDim(debug_info, 1));
+
+  TFE_DeleteTensorDebugInfo(debug_info);
+  TFE_DeleteTensorHandle(h);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+}

tensorflow/c/eager/c_api_distributed_test.cc

@@ -0,0 +1,627 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <regex>  // NOLINT
+#include <string>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/common_runtime/function_optimization_registry.h"
+#include "tensorflow/core/common_runtime/optimization_registry.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+namespace {
+
+using ::tensorflow::string;
+
+// Add the values of three variables on three different tasks.
+string AddVariablesFunction() {
+  tensorflow::FunctionDef def;
+  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(
+      "    signature {"
+      "      name: 'AddVariablesFunction'"
+      "      input_arg {"
+      "        name: 'var'"
+      "        type: DT_RESOURCE"
+      "      }"
+      "      output_arg {"
+      "        name: 'sum'"
+      "        type: DT_FLOAT"
+      "      }"
+      "    }"
+      "    node_def {"
+      "      name: 'read0'"
+      "      op: 'ReadVariableOp'"
+      "      input: 'var'"
+      "      device: '/job:localhost/replica:0/task:0/device:CPU:0'"
+      "      attr {"
+      "        key: 'dtype'"
+      "        value {"
+      "          type: DT_FLOAT"
+      "        }"
+      "      }"
+      "    }"
+      "    node_def {"
+      "      name: 'read1'"
+      "      op: 'ReadVariableOp'"
+      "      input: 'var'"
+      "      device: '/job:localhost/replica:0/task:1/device:CPU:0'"
+      "      attr {"
+      "        key: 'dtype'"
+      "        value {"
+      "          type: DT_FLOAT"
+      "        }"
+      "      }"
+      "    }"
+      "    node_def {"
+      "      name: 'read2'"
+      "      op: 'ReadVariableOp'"
+      "      input: 'var'"
+      "      device: '/job:localhost/replica:0/task:2/device:CPU:0'"
+      "      attr {"
+      "        key: 'dtype'"
+      "        value {"
+      "          type: DT_FLOAT"
+      "        }"
+      "      }"
+      "    }"
+      "    node_def {"
+      "      name: 'add1'"
+      "      op: 'Add'"
+      "      input: 'read0:value:0'"
+      "      input: 'read1:value:0'"
+      "      attr {"
+      "        key: 'T'"
+      "        value {"
+      "          type: DT_FLOAT"
+      "        }"
+      "      }"
+      "    }"
+      "    node_def {"
+      "      name: 'add2'"
+      "      op: 'Add'"
+      "      input: 'add1:z:0'"
+      "      input: 'read2:value:0'"
+      "      attr {"
+      "        key: 'T'"
+      "        value {"
+      "          type: DT_FLOAT"
+      "        }"
+      "      }"
+      "    }"
+      "    ret {"
+      "      key: 'sum'"
+      "      value: 'add2:z:0'"
+      "    }",
+      &def));
+  return def.SerializeAsString();
+}
+
+void TestFunctionWithPackedInput(const bool remote) {
+  tensorflow::ServerDef server_def = GetServerDef(3);
+
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(/*enable=*/true));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  const char task0_name[] = "/job:localhost/replica:0/task:0/device:CPU:0";
+  const char task1_name[] = "/job:localhost/replica:0/task:1/device:CPU:0";
+  const char task2_name[] = "/job:localhost/replica:0/task:2/device:CPU:0";
+
+  // Create one variable per task.
+  TFE_TensorHandle* h0 = TestVariable(ctx, 1.0, task1_name);
+  TFE_TensorHandle* h1 = TestVariable(ctx, 2.0, task2_name);
+  TFE_TensorHandle* h2 = TestVariable(ctx, 3.0, task0_name);
+
+  // Add a sync point to make sure that variables have been initialized
+  // before the function execution starts.
+  TFE_ContextAsyncWait(ctx, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // Pack 3 variable handles into one TFE_TensorHandle.
+  // When remote is false, function device is placed on task0. Handle types are
+  // REMOTE, REMOTE, LOCAL on task0. When remote is true, function device is
+  // placed on task1, Handle types are LOCAL, REMOTE, LOCAL on task1.
+  int num_replicas = 3;
+  std::vector<TFE_TensorHandle*> handles = {h0, h1, h2};
+  TFE_TensorHandle* packed_handle =
+      TFE_CreatePackedTensorHandle(ctx, handles.data(), &num_replicas, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  EXPECT_EQ(TFE_TensorHandleDataType(packed_handle), TF_RESOURCE);
+  EXPECT_EQ(TFE_TensorHandleNumDims(packed_handle, status), 0);
+  EXPECT_EQ(TFE_TensorHandleNumElements(packed_handle, status), 1);
+
+  const string composite_device_name =
+      "/job:localhost/replica:0/task:0/device:COMPOSITE:0";
+  EXPECT_EQ(TFE_TensorHandleDeviceName(packed_handle, status),
+            composite_device_name);
+  EXPECT_EQ(TFE_TensorHandleBackingDeviceName(packed_handle, status),
+            composite_device_name);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  // Register and run a function which returns the sum of 3 variables.
+  const string function_def = AddVariablesFunction();
+  TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(),
+                            status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_Op* func = TFE_NewOp(ctx, "AddVariablesFunction", status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_OpAddInput(func, packed_handle, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  if (remote) {
+    TFE_OpSetDevice(func, task1_name, status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  }
+
+  TFE_TensorHandle* retvals[1] = {nullptr};
+  int num_retvals = 1;
+  TFE_Execute(func, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  ASSERT_EQ(1, num_retvals);
+  TFE_DeleteOp(func);
+  TFE_DeleteTensorHandle(packed_handle);
+  TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteTensorHandle(retvals[0]);
+  float sum = 0;
+  EXPECT_EQ(sizeof(sum), TF_TensorByteSize(t));
+  memcpy(&sum, TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+  EXPECT_EQ(sum, 6.0);
+
+  TFE_DeleteTensorHandle(h0);
+  TFE_DeleteTensorHandle(h1);
+  TFE_DeleteTensorHandle(h2);
+
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+  TFE_ContextRemoveFunction(ctx, "AddVariablesFunction", status);
+  TFE_DeleteContext(ctx);
+
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server1.release();
+  worker_server2.release();
+}
+
+TEST(CAPI, TestLocalFunctionWithPackedInput) {
+  TestFunctionWithPackedInput(/*remote=*/false);
+}
+
+TEST(CAPI, TestRemoteFunctionWithPackedInput) {
+  TestFunctionWithPackedInput(/*remote=*/true);
+}
+
+string VariableAddFunctionSignature() {
+  return "    signature {"
+         "      name: 'VariableAddFunction'"
+         "      input_arg {"
+         "        name: 'var0'"
+         "        type: DT_RESOURCE"
+         "      }"
+         "      output_arg {"
+         "        name: 'var0_value'"
+         "        type: DT_FLOAT"
+         "      }"
+         "    }"
+         "    node_def {"
+         "      name: 'read0'"
+         "      op: 'ReadVariableOp'"
+         "      input: 'var0'"
+         "      attr {"
+         "        key: 'dtype'"
+         "        value {"
+         "          type: DT_FLOAT"
+         "        }"
+         "      }"
+         "    }"
+         "    node_def {"
+         "      name: 'add'"
+         "      op: 'Add'"
+         "      input: 'read0:value:0'"
+         "      input: 'read0:value:0'"
+         "      device: '/job:localhost/task:1/device:CPU:0'"
+         "      attr {"
+         "        key: 'T'"
+         "        value {"
+         "          type: DT_FLOAT"
+         "        }"
+         "      }"
+         "    }"
+         "    node_def {"
+         "      name: 'identity'"
+         "      op: 'Identity'"
+         "      input: 'add:z:0'"
+         "      device: '/job:localhost/task:0/device:CPU:0'"
+         "      attr {"
+         "        key: 'T'"
+         "        value {"
+         "          type: DT_FLOAT"
+         "        }"
+         "      }"
+         "    }"
+         "    ret {"
+         "      key: 'var0_value'"
+         "      value: 'identity:output:0'"
+         "    }";
+}
+
+string VariableAddFunction() {
+  tensorflow::FunctionDef def;
+  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(
+      VariableAddFunctionSignature(), &def));
+  return def.SerializeAsString();
+}
+
+// A graph optimization pass that would fail when triggered for more than once.
+class GraphErrorInjectionPass : public tensorflow::GraphOptimizationPass {
+ public:
+  static bool enabled_;
+  GraphErrorInjectionPass() {}
+
+  absl::Status Run(
+      const tensorflow::GraphOptimizationPassOptions& options) override {
+    if (!enabled_) {
+      return absl::OkStatus();
+    }
+    if (first_call_) {
+      first_call_ = false;
+      return absl::OkStatus();
+    }
+    return absl::InternalError("Graph pass runs for more than once!");
+  }
+
+ private:
+  bool first_call_ = true;
+};
+
+// After the graph pass is registered, it takes effect globally and can affect
+// other test cases. Define a static variable to switch it on and off.
+bool GraphErrorInjectionPass::enabled_ = false;
+
+// Test to ensure that a registered graph optimization pass is only executed
+// once (i.e., on the main function side) in running distributed functions.
+// This test creates a cluster with two workers, create a variable on the
+// second worker, and run a distributed function (VariableAddFunction) whose ops
+// span the local and remote workers. If the graph optimization pass is executed
+// on both the main function side and the component function side, an error will
+// be thrown in the registered graph optimization pass.
+TEST(CAPI, DistributedFunctionGraphPassOnlyOnce) {
+  // Register graph pass that will raise error if called more than once.
+  tensorflow::optimization_registration::OptimizationPassRegistration
+      register_test_pass(tensorflow::OptimizationPassRegistry::PRE_PLACEMENT, 0,
+                         std::make_unique<GraphErrorInjectionPass>(),
+                         "error_injector");
+  GraphErrorInjectionPass::enabled_ = true;
+
+  tensorflow::ServerDef server_def = GetServerDef(3);
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+  const char dev2_name[] = "/job:localhost/replica:0/task:2/device:CPU:0";
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* var_handle = TestVariable(ctx, 2.0, dev2_name);
+  EXPECT_NE(var_handle, nullptr);
+  TFE_ContextAsyncWait(ctx, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  const string function_def = VariableAddFunction();
+  TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(),
+                            status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_Op* func = TFE_NewOp(ctx, "VariableAddFunction", status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_OpAddInput(func, var_handle, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_TensorHandle* retvals[1] = {nullptr};
+  int num_retvals = 1;
+  TFE_Execute(func, &retvals[0], &num_retvals, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  ASSERT_EQ(1, num_retvals);
+  TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteTensorHandle(retvals[0]);
+  float sum = 0;
+  ASSERT_EQ(sizeof(sum), TF_TensorByteSize(t));
+  memcpy(&sum, TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+  ASSERT_EQ(sum, 4.0);
+
+  TFE_DeleteOp(func);
+  TFE_DeleteTensorHandle(var_handle);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server1.release();
+  worker_server2.release();
+
+  // Disable the test graph pass so it does not affect other test cases.
+  GraphErrorInjectionPass::enabled_ = false;
+}
+
+string VariableAddFunctionWithGraphError() {
+  string signature = VariableAddFunctionSignature();
+  // Replace the node 'read0' with 'read0_maybe_with_graph_error', so that the
+  // error injecting pass can identify and introduce graph pass errors.
+  signature = std::regex_replace(signature, std::regex("read0"),
+                                 "read0_maybe_with_graph_error");
+  tensorflow::FunctionDef def;
+  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(signature, &def));
+  return def.SerializeAsString();
+}
+
+class FunctionErrorInjectionPass : public tensorflow::FunctionOptimizationPass {
+ public:
+  FunctionErrorInjectionPass(string error_node, string error_device)
+      : error_node_(error_node), error_device_(error_device) {}
+  absl::Status Run(const std::string& function_name,
+                   const tensorflow::DeviceSet& device_set,
+                   const tensorflow::ConfigProto& config_proto,
+                   const FunctionOptions& function_options,
+                   std::unique_ptr<tensorflow::Graph>* graph,
+                   tensorflow::FunctionLibraryDefinition* flib_def,
+                   std::vector<std::string>* control_ret_node_names,
+                   bool* control_rets_updated) override {
+    // Inject failure to function instantiation if finding a node that contains
+    // the given node name (error_node_) and requested device (error_device_).
+    for (const auto node : graph->get()->nodes()) {
+      if (node->name().find(error_node_) != string::npos &&
+          node->requested_device() == error_device_) {
+        return absl::InternalError("Injected graph pass error.");
+      }
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  const string error_node_;
+  const string error_device_;
+};
+
+void TestDistributedFunctionCancellation(bool inject_error) {
+  tensorflow::ServerDef server_def = GetServerDef(3);
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+  const char dev2_name[] = "/job:localhost/replica:0/task:2/device:CPU:0";
+
+  if (inject_error) {
+    // Inject a function optimization pass failure when it sees the
+    // 'read0_maybe_with_graph_error' op having a requested device `dev2_name`.
+    // During execution:
+    //   * task:0 processes main function `VariableAddFunctionWithGraphError`
+    //     and places the 'read0_maybe_with_graph_error' op on task:2
+    //   * task:0 partitions the main function with a subgraph containing
+    //     'read0_maybe_with_graph_error' sent to task:2
+    //   * task:2 graph pass reports an error when it sees
+    //     'read0_maybe_with_graph_error' with dev2_name
+    tensorflow::function_optimization_registration::
+        FunctionOptimizationPassRegistration register_test_pass(
+            std::make_unique<FunctionErrorInjectionPass>(
+                "read0_maybe_with_graph_error", dev2_name));
+  }
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* var_handle = TestVariable(ctx, 2.0, dev2_name);
+  EXPECT_NE(var_handle, nullptr);
+  TFE_ContextAsyncWait(ctx, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  const string function_def = inject_error ? VariableAddFunctionWithGraphError()
+                                           : VariableAddFunction();
+  TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(),
+                            status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_Op* func = TFE_NewOp(ctx, "VariableAddFunction", status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_OpAddInput(func, var_handle, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_TensorHandle* retvals[1] = {nullptr};
+  int num_retvals = 1;
+  TFE_Execute(func, &retvals[0], &num_retvals, status);
+
+  if (inject_error) {
+    ASSERT_EQ(TF_INTERNAL, TF_GetCode(status)) << TF_Message(status);
+  } else {
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    ASSERT_EQ(1, num_retvals);
+    TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_DeleteTensorHandle(retvals[0]);
+    float sum = 0;
+    ASSERT_EQ(sizeof(sum), TF_TensorByteSize(t));
+    memcpy(&sum, TF_TensorData(t), TF_TensorByteSize(t));
+    TF_DeleteTensor(t);
+    ASSERT_EQ(sum, 4.0);
+  }
+
+  TFE_DeleteOp(func);
+  TFE_DeleteTensorHandle(var_handle);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server1.release();
+  worker_server2.release();
+}
+
+TEST(CAPI, DistributedFunctionNoError) {
+  TestDistributedFunctionCancellation(false);
+}
+
+// TODO(b/170399182): Update test once an alternative to using the function
+// optimization hook is in place.
+TEST(CAPI, DISABLED_DistributedFunctionCancelledOnError) {
+  TestDistributedFunctionCancellation(true);
+}
+
+void TestRemoteExecuteDeleteContextWithOutstandingRPC(bool async) {
+  tensorflow::ServerDef server_def = GetServerDef(2);
+
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts,
+                                             TFE_DEVICE_PLACEMENT_EXPLICIT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  // Use large matrices so that RPCs don't return before we get a chance
+  // to call TFE_DeleteContext.
+  TFE_TensorHandle* h0_task0 = TestMatrixTensorHandle100x100(ctx);
+  TFE_TensorHandle* h1_task0 = TestMatrixTensorHandle100x100(ctx);
+  const char remote_device_name[] =
+      "/job:localhost/replica:0/task:1/device:CPU:0";
+  auto* h0_task1 =
+      TFE_TensorHandleCopyToDevice(h0_task0, ctx, remote_device_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* h1_task1 =
+      TFE_TensorHandleCopyToDevice(h1_task0, ctx, remote_device_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
+  TFE_OpSetDevice(matmul, remote_device_name, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* retvals[1];
+  int num_retvals = 1;
+  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+
+  TFE_DeleteTensorHandle(h0_task0);
+  TFE_DeleteTensorHandle(h1_task0);
+  TFE_DeleteTensorHandle(h0_task1);
+  TFE_DeleteTensorHandle(h1_task1);
+  TFE_DeleteTensorHandle(retvals[0]);
+
+  TFE_DeleteOp(matmul);
+
+  TFE_DeleteContext(ctx);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+}
+
+TEST(CAPI, RemoteExecuteDeleteContextWithOutstandingRPC) {
+  TestRemoteExecuteDeleteContextWithOutstandingRPC(false);
+}
+
+TEST(CAPI, RemoteExecuteDeleteContextWithOutstandingRPCAsync) {
+  TestRemoteExecuteDeleteContextWithOutstandingRPC(true);
+}
+}  // namespace

tensorflow/c/eager/c_api_experimental.cc

@@ -0,0 +1,993 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_experimental.h"
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/match.h"
+#include "absl/time/time.h"
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/eager/tfe_op_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "xla/tsl/c/tsl_status_internal.h"
+#include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
+#include "xla/tsl/framework/cancellation.h"
+#include "tensorflow/core/common_runtime/composite_device.h"
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/distributed_runtime/coordination/coordination_service_error_util.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+#include "tensorflow/core/lib/monitoring/gauge.h"
+#include "tensorflow/core/lib/monitoring/sampler.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/strcat.h"
+
+using tensorflow::string;
+
+void TFE_OpReset(TFE_Op* op_to_reset, const char* op_or_function_name,
+                 const char* raw_device_name, TF_Status* status) {
+  if (op_to_reset) {
+    tensorflow::ImmediateExecutionOperation* op =
+        tensorflow::unwrap(op_to_reset);
+    op->Clear();
+    status->status = op->Reset(op_or_function_name, raw_device_name);
+  } else {
+    TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                 "op_to_reset should not be nullptr");
+  }
+}
+
+void TFE_ContextEnableGraphCollection(TFE_Context* ctx) {
+  tensorflow::unwrap(ctx)->SetShouldStoreGraphs(true);
+}
+
+void TFE_ContextDisableGraphCollection(TFE_Context* ctx) {
+  tensorflow::unwrap(ctx)->SetShouldStoreGraphs(false);
+}
+
+uint64_t TFE_GetContextId(TFE_Context* ctx) {
+  tensorflow::EagerContext* context =
+      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
+  return context->GetContextId();
+}
+
+void TFE_MonitoringCounterCellIncrementBy(TFE_MonitoringCounterCell* cell,
+                                          int64_t value) {
+  cell->cell.IncrementBy(value);
+}
+
+int64_t TFE_MonitoringCounterCellValue(TFE_MonitoringCounterCell* cell) {
+  return cell->cell.value();
+}
+
+TFE_MonitoringCounter0* TFE_MonitoringNewCounter0(const char* name,
+                                                  TF_Status* status,
+                                                  const char* description) {
+  auto* result = new TFE_MonitoringCounter0({name, description});
+  tsl::Set_TF_Status_from_Status(status, result->counter->GetStatus());
+  if (!result->counter->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteCounter0(TFE_MonitoringCounter0* counter) {
+  delete counter;
+}
+
+TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter0(
+    TFE_MonitoringCounter0* counter) {
+  return static_cast<TFE_MonitoringCounterCell*>(
+      static_cast<void*>(counter->counter->GetCell()));
+}
+
+TFE_MonitoringCounter1* TFE_MonitoringNewCounter1(const char* name,
+                                                  TF_Status* status,
+                                                  const char* description,
+                                                  const char* label1) {
+  auto* result = new TFE_MonitoringCounter1({name, description, label1});
+  tsl::Set_TF_Status_from_Status(status, result->counter->GetStatus());
+  if (!result->counter->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteCounter1(TFE_MonitoringCounter1* counter) {
+  delete counter;
+}
+
+TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter1(
+    TFE_MonitoringCounter1* counter, const char* label1) {
+  return static_cast<TFE_MonitoringCounterCell*>(
+      static_cast<void*>(counter->counter->GetCell(label1)));
+}
+
+TFE_MonitoringCounter2* TFE_MonitoringNewCounter2(const char* name,
+                                                  TF_Status* status,
+                                                  const char* description,
+                                                  const char* label1,
+                                                  const char* label2) {
+  auto* result =
+      new TFE_MonitoringCounter2({name, description, label1, label2});
+  tsl::Set_TF_Status_from_Status(status, result->counter->GetStatus());
+  if (!result->counter->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteCounter2(TFE_MonitoringCounter2* counter) {
+  delete counter;
+}
+
+TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter2(
+    TFE_MonitoringCounter2* counter, const char* label1, const char* label2) {
+  return static_cast<TFE_MonitoringCounterCell*>(
+      static_cast<void*>(counter->counter->GetCell(label1, label2)));
+}
+
+void TFE_MonitoringIntGaugeCellSet(TFE_MonitoringIntGaugeCell* cell,
+                                   int64_t value) {
+  cell->cell.Set(value);
+}
+
+int64_t TFE_MonitoringIntGaugeCellValue(TFE_MonitoringIntGaugeCell* cell) {
+  return cell->cell.value();
+}
+
+TFE_MonitoringIntGauge0* TFE_MonitoringNewIntGauge0(const char* name,
+                                                    TF_Status* status,
+                                                    const char* description) {
+  auto* result = new TFE_MonitoringIntGauge0({name, description});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteIntGauge0(TFE_MonitoringIntGauge0* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringIntGaugeCell* TFE_MonitoringGetCellIntGauge0(
+    TFE_MonitoringIntGauge0* gauge) {
+  return static_cast<TFE_MonitoringIntGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell()));
+}
+
+TFE_MonitoringIntGauge1* TFE_MonitoringNewIntGauge1(const char* name,
+                                                    TF_Status* status,
+                                                    const char* description,
+                                                    const char* label1) {
+  auto* result = new TFE_MonitoringIntGauge1({name, description, label1});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteIntGauge1(TFE_MonitoringIntGauge1* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringIntGaugeCell* TFE_MonitoringGetCellIntGauge1(
+    TFE_MonitoringIntGauge1* gauge, const char* label1) {
+  return static_cast<TFE_MonitoringIntGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1)));
+}
+
+TFE_MonitoringIntGauge2* TFE_MonitoringNewIntGauge2(const char* name,
+                                                    TF_Status* status,
+                                                    const char* description,
+                                                    const char* label1,
+                                                    const char* label2) {
+  auto* result =
+      new TFE_MonitoringIntGauge2({name, description, label1, label2});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteIntGauge2(TFE_MonitoringIntGauge2* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringIntGaugeCell* TFE_MonitoringGetCellIntGauge2(
+    TFE_MonitoringIntGauge2* gauge, const char* label1, const char* label2) {
+  return static_cast<TFE_MonitoringIntGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1, label2)));
+}
+
+void TFE_MonitoringStringGaugeCellSet(TFE_MonitoringStringGaugeCell* cell,
+                                      const char* value) {
+  cell->cell.Set({value});
+}
+
+const void TFE_MonitoringStringGaugeCellValue(
+    TFE_MonitoringStringGaugeCell* cell, TF_Buffer* buf) {
+  std::string value = cell->cell.value();
+  void* data = tensorflow::port::Malloc(value.length());
+  value.copy(static_cast<char*>(data), value.length(), 0);
+  buf->data = data;
+  buf->length = value.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+}
+
+TFE_MonitoringStringGauge0* TFE_MonitoringNewStringGauge0(
+    const char* name, TF_Status* status, const char* description) {
+  auto* result = new TFE_MonitoringStringGauge0({name, description});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteStringGauge0(TFE_MonitoringStringGauge0* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringStringGaugeCell* TFE_MonitoringGetCellStringGauge0(
+    TFE_MonitoringStringGauge0* gauge) {
+  return static_cast<TFE_MonitoringStringGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell()));
+}
+
+TFE_MonitoringStringGauge1* TFE_MonitoringNewStringGauge1(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1) {
+  auto* result = new TFE_MonitoringStringGauge1({name, description, label1});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteStringGauge1(TFE_MonitoringStringGauge1* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringStringGaugeCell* TFE_MonitoringGetCellStringGauge1(
+    TFE_MonitoringStringGauge1* gauge, const char* label1) {
+  return static_cast<TFE_MonitoringStringGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1)));
+}
+
+TFE_MonitoringStringGauge2* TFE_MonitoringNewStringGauge2(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1, const char* label2) {
+  auto* result =
+      new TFE_MonitoringStringGauge2({name, description, label1, label2});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteStringGauge2(TFE_MonitoringStringGauge2* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringStringGaugeCell* TFE_MonitoringGetCellStringGauge2(
+    TFE_MonitoringStringGauge2* gauge, const char* label1, const char* label2) {
+  return static_cast<TFE_MonitoringStringGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1, label2)));
+}
+
+TFE_MonitoringStringGauge3* TFE_MonitoringNewStringGauge3(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1, const char* label2, const char* label3) {
+  auto* result = new TFE_MonitoringStringGauge3(
+      {name, description, label1, label2, label3});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteStringGauge3(TFE_MonitoringStringGauge3* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringStringGaugeCell* TFE_MonitoringGetCellStringGauge3(
+    TFE_MonitoringStringGauge3* gauge, const char* label1, const char* label2,
+    const char* label3) {
+  return static_cast<TFE_MonitoringStringGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1, label2, label3)));
+}
+
+TFE_MonitoringStringGauge4* TFE_MonitoringNewStringGauge4(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1, const char* label2, const char* label3,
+    const char* label4) {
+  auto* result = new TFE_MonitoringStringGauge4(
+      {name, description, label1, label2, label3, label4});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteStringGauge4(TFE_MonitoringStringGauge4* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringStringGaugeCell* TFE_MonitoringGetCellStringGauge4(
+    TFE_MonitoringStringGauge4* gauge, const char* label1, const char* label2,
+    const char* label3, const char* label4) {
+  return static_cast<TFE_MonitoringStringGaugeCell*>(static_cast<void*>(
+      gauge->gauge->GetCell(label1, label2, label3, label4)));
+}
+
+void TFE_MonitoringBoolGaugeCellSet(TFE_MonitoringBoolGaugeCell* cell,
+                                    bool value) {
+  cell->cell.Set(value);
+}
+
+bool TFE_MonitoringBoolGaugeCellValue(TFE_MonitoringBoolGaugeCell* cell) {
+  return cell->cell.value();
+}
+
+TFE_MonitoringBoolGauge0* TFE_MonitoringNewBoolGauge0(const char* name,
+                                                      TF_Status* status,
+                                                      const char* description) {
+  auto* result = new TFE_MonitoringBoolGauge0({name, description});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteBoolGauge0(TFE_MonitoringBoolGauge0* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringBoolGaugeCell* TFE_MonitoringGetCellBoolGauge0(
+    TFE_MonitoringBoolGauge0* gauge) {
+  return static_cast<TFE_MonitoringBoolGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell()));
+}
+
+TFE_MonitoringBoolGauge1* TFE_MonitoringNewBoolGauge1(const char* name,
+                                                      TF_Status* status,
+                                                      const char* description,
+                                                      const char* label1) {
+  auto* result = new TFE_MonitoringBoolGauge1({name, description, label1});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteBoolGauge1(TFE_MonitoringBoolGauge1* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringBoolGaugeCell* TFE_MonitoringGetCellBoolGauge1(
+    TFE_MonitoringBoolGauge1* gauge, const char* label1) {
+  return static_cast<TFE_MonitoringBoolGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1)));
+}
+
+TFE_MonitoringBoolGauge2* TFE_MonitoringNewBoolGauge2(const char* name,
+                                                      TF_Status* status,
+                                                      const char* description,
+                                                      const char* label1,
+                                                      const char* label2) {
+  auto* result =
+      new TFE_MonitoringBoolGauge2({name, description, label1, label2});
+  tsl::Set_TF_Status_from_Status(status, result->gauge->GetStatus());
+  if (!result->gauge->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteBoolGauge2(TFE_MonitoringBoolGauge2* gauge) {
+  delete gauge;
+}
+
+TFE_MonitoringBoolGaugeCell* TFE_MonitoringGetCellBoolGauge2(
+    TFE_MonitoringBoolGauge2* gauge, const char* label1, const char* label2) {
+  return static_cast<TFE_MonitoringBoolGaugeCell*>(
+      static_cast<void*>(gauge->gauge->GetCell(label1, label2)));
+}
+
+void TFE_MonitoringSamplerCellAdd(TFE_MonitoringSamplerCell* cell,
+                                  double value) {
+  cell->cell.Add(value);
+}
+
+void TFE_MonitoringSamplerCellValue(TFE_MonitoringSamplerCell* cell,
+                                    TF_Buffer* buf) {
+  string content;
+  cell->cell.value().SerializeToString(&content);
+  void* data = tensorflow::port::Malloc(content.length());
+  content.copy(static_cast<char*>(data), content.length(), 0);
+  buf->data = data;
+  buf->length = content.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+}
+
+TFE_MonitoringBuckets* TFE_MonitoringNewExponentialBuckets(double scale,
+                                                           double growth_factor,
+                                                           int bucket_count) {
+  return new TFE_MonitoringBuckets([scale, growth_factor, bucket_count]() {
+    return tensorflow::monitoring::Buckets::Exponential(scale, growth_factor,
+                                                        bucket_count);
+  });
+}
+
+void TFE_MonitoringDeleteBuckets(TFE_MonitoringBuckets* buckets) {
+  delete buckets;
+}
+
+TFE_MonitoringSampler0* TFE_MonitoringNewSampler0(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* status,
+    const char* description) {
+  auto* result = new TFE_MonitoringSampler0(
+      {name, buckets->create_buckets(), description});
+  tsl::Set_TF_Status_from_Status(status, result->sampler->GetStatus());
+  if (!result->sampler->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteSampler0(TFE_MonitoringSampler0* sampler) {
+  delete sampler;
+}
+
+TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler0(
+    TFE_MonitoringSampler0* sampler) {
+  return static_cast<TFE_MonitoringSamplerCell*>(
+      static_cast<void*>(sampler->sampler->GetCell()));
+}
+
+TFE_MonitoringSampler1* TFE_MonitoringNewSampler1(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* status,
+    const char* description, const char* label1) {
+  auto* result = new TFE_MonitoringSampler1(
+      {name, buckets->create_buckets(), description, label1});
+  tsl::Set_TF_Status_from_Status(status, result->sampler->GetStatus());
+  if (!result->sampler->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteSampler1(TFE_MonitoringSampler1* sampler) {
+  delete sampler;
+}
+
+TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler1(
+    TFE_MonitoringSampler1* sampler, const char* label1) {
+  return static_cast<TFE_MonitoringSamplerCell*>(
+      static_cast<void*>(sampler->sampler->GetCell(label1)));
+}
+
+TFE_MonitoringSampler2* TFE_MonitoringNewSampler2(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* status,
+    const char* description, const char* label1, const char* label2) {
+  auto* result = new TFE_MonitoringSampler2(
+      {name, buckets->create_buckets(), description, label1, label2});
+  tsl::Set_TF_Status_from_Status(status, result->sampler->GetStatus());
+  if (!result->sampler->GetStatus().ok()) {
+    delete result;
+    return nullptr;
+  }
+  return result;
+}
+
+void TFE_MonitoringDeleteSampler2(TFE_MonitoringSampler2* sampler) {
+  delete sampler;
+}
+
+TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler2(
+    TFE_MonitoringSampler2* sampler, const char* label1, const char* label2) {
+  return static_cast<TFE_MonitoringSamplerCell*>(
+      static_cast<void*>(sampler->sampler->GetCell(label1, label2)));
+}
+
+void TFE_ContextOptionsSetTfrt(TFE_ContextOptions* options, bool use_tfrt) {
+  options->use_tfrt = use_tfrt;
+}
+
+TFE_CancellationManager* TFE_NewCancellationManager() {
+  return tensorflow::wrap(new tensorflow::CancellationManager);
+}
+
+void TFE_CancellationManagerStartCancel(
+    TFE_CancellationManager* cancellation_manager) {
+  tensorflow::unwrap(cancellation_manager)->StartCancel();
+}
+
+bool TFE_CancellationManagerIsCancelled(
+    TFE_CancellationManager* cancellation_manager) {
+  return tensorflow::unwrap(cancellation_manager)->IsCancelled();
+}
+
+bool TFE_CancellationManagerIsCancelling(
+    TFE_CancellationManager* cancellation_manager) {
+  return tensorflow::unwrap(cancellation_manager)->IsCancelling();
+}
+
+TFE_CancellationToken TFE_CancellationManagerGetToken(
+    TFE_CancellationManager* cancellation_manager) {
+  return tensorflow::unwrap(cancellation_manager)->get_cancellation_token();
+}
+
+bool TFE_CancellationManagerRegisterCallback(
+    TFE_CancellationManager* cancellation_manager, TFE_CancellationToken token,
+    const TFE_CancelCallback* c_callback, const char* callback_name) {
+  tensorflow::CancelCallback callback = [callback = c_callback->callback,
+                                         context = c_callback->context]() {
+    callback(context);
+  };
+  return tensorflow::unwrap(cancellation_manager)
+      ->RegisterCallbackWithErrorLogging(token, callback, callback_name);
+}
+
+bool TFE_CancellationManagerDeregisterCallback(
+    TFE_CancellationManager* cancellation_manager,
+    TFE_CancellationToken token) {
+  return tensorflow::unwrap(cancellation_manager)->DeregisterCallback(token);
+}
+
+bool TFE_CancellationManagerTryDeregisterCallback(
+    TFE_CancellationManager* cancellation_manager,
+    TFE_CancellationToken token) {
+  return tensorflow::unwrap(cancellation_manager)->TryDeregisterCallback(token);
+}
+
+void TFE_DeleteCancellationManager(
+    TFE_CancellationManager* cancellation_manager) {
+  delete tensorflow::unwrap(cancellation_manager);
+}
+
+void TFE_OpSetCancellationManager(TFE_Op* op,
+                                  TFE_CancellationManager* cancellation_manager,
+                                  TF_Status* status) {
+  tensorflow::unwrap(op)->SetCancellationManager(
+      tensorflow::unwrap(cancellation_manager));
+  status->status = absl::OkStatus();
+}
+
+TFE_Executor* TFE_NewExecutor(bool is_async, bool enable_streaming_enqueue,
+                              int in_flight_nodes_limit) {
+  return new TFE_Executor(is_async, enable_streaming_enqueue,
+                          in_flight_nodes_limit);
+}
+
+void TFE_DeleteExecutor(TFE_Executor* executor) { delete executor; }
+
+bool TFE_ExecutorIsAsync(TFE_Executor* executor) {
+  return executor->executor()->Async();
+}
+
+void TFE_ExecutorWaitForAllPendingNodes(TFE_Executor* executor,
+                                        TF_Status* status) {
+  status->status = executor->executor()->WaitForAllPendingNodes();
+}
+
+void TFE_ExecutorClearError(TFE_Executor* executor) {
+  executor->executor()->ClearError();
+}
+
+void TFE_ContextSetExecutorForThread(TFE_Context* ctx, TFE_Executor* executor) {
+  tensorflow::unwrap(ctx)->SetExecutorForThread(executor->executor());
+}
+
+TFE_Executor* TFE_ContextGetExecutorForThread(TFE_Context* ctx) {
+  return new TFE_Executor(&tensorflow::unwrap(ctx)->Executor());
+}
+
+void TFE_HostAddressSpace(TFE_Context* ctx, TF_Buffer* buf) {
+  auto address_space = tensorflow::DeviceNameUtils::AddressSpace(
+      tensorflow::unwrap(ctx)->HostCPUParsedName());
+  auto str = tensorflow::DeviceNameUtils::ParsedNameToString(address_space);
+  void* data = tensorflow::port::Malloc(str.length());
+  str.copy(static_cast<char*>(data), str.length(), 0);
+  buf->data = data;
+  buf->length = str.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+}
+
+void TFE_ContextGetFunctionDef(TFE_Context* ctx, const char* function_name,
+                               TF_Buffer* buf, TF_Status* status) {
+  auto* function_def = tensorflow::unwrap(ctx)->FindFunctionDef(function_name);
+  if (function_def == nullptr) {
+    status->status = absl::NotFoundError(
+        absl::StrCat("Unable to find FunctionDef with name: ", function_name));
+    return;
+  }
+  string str = function_def->SerializeAsString();
+  void* data = tensorflow::port::Malloc(str.length());
+  str.copy(static_cast<char*>(data), str.length(), 0);
+  buf->data = data;
+  buf->length = str.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+  status->status = absl::OkStatus();
+}
+
+void TFE_ContextGetGraphDebugInfo(TFE_Context* ctx, const char* function_name,
+                                  TF_Buffer* buf, TF_Status* status) {
+  auto function_record = tensorflow::unwrap(ctx)->FindRecord(function_name);
+  if (function_record == nullptr) {
+    status->status = absl::NotFoundError(
+        absl::StrCat("Unable to find function with name: ", function_name));
+    return;
+  }
+
+  tensorflow::GraphDebugInfo debug_info =
+      tensorflow::StackTracesMapToGraphDebugInfo(
+          function_record->stack_traces());
+
+  string str = debug_info.SerializeAsString();
+  void* data = tensorflow::port::Malloc(str.length());
+  str.copy(static_cast<char*>(data), str.length(), 0);
+  buf->data = data;
+  buf->length = str.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+  status->status = absl::OkStatus();
+}
+
+TF_Tensor* TFE_AllocateHostTensor(TFE_Context* ctx, TF_DataType dtype,
+                                  const int64_t* dims, int num_dims,
+                                  TF_Status* status) {
+  std::vector<int64_t> dimvec(num_dims);
+  for (int i = 0; i < num_dims; ++i) {
+    dimvec[i] = static_cast<int64_t>(dims[i]);
+  }
+
+  if (ctx == nullptr) {
+    status->status = absl::InvalidArgumentError("Invalid Context");
+    return nullptr;
+  }
+
+  tensorflow::AbstractTensorInterface* t =
+      tensorflow::unwrap(ctx)->CreateTensor(
+          static_cast<tensorflow::DataType>(dtype), dimvec);
+
+  if (t == nullptr) {
+    status->status =
+        absl::InvalidArgumentError(absl::StrCat("Unsupported dtype: ", dtype));
+    return nullptr;
+  }
+
+  return new TF_Tensor{t};
+}
+
+TFE_TensorHandle* TFE_NewTensorHandleFromTensor(TFE_Context* ctx, TF_Tensor* t,
+                                                TF_Status* status) {
+  return tensorflow::wrap(
+      tensorflow::unwrap(ctx)->CreateLocalHandle(t->tensor));
+}
+
+TFE_TensorHandle* TFE_CreatePackedTensorHandle(TFE_Context* ctx,
+                                               TFE_TensorHandle** handles,
+                                               int* num_handles,
+                                               TF_Status* status) {
+  std::vector<tensorflow::TensorHandle*> tensor_handles;
+  tensor_handles.reserve(*num_handles);
+  for (int i = 0; i < *num_handles; ++i) {
+    tensorflow::ImmediateExecutionTensorHandle* unwrapped_handle =
+        tensorflow::unwrap(handles[i]);
+    if (tensorflow::CustomDeviceTensorHandle::classof(unwrapped_handle)) {
+      // One of the inputs we're trying to pack is on a custom device. We'll let
+      // the first custom device we see handle all of the packing.
+      auto* custom_device_handle =
+          absl::down_cast<tensorflow::CustomDeviceTensorHandle*>(
+              unwrapped_handle);
+      tensorflow::ImmediateExecutionTensorHandle* result;
+      status->status = custom_device_handle->device()->Pack(
+          absl::Span<tensorflow::ImmediateExecutionTensorHandle*>(
+              tensorflow::unwrap(handles), *num_handles),
+          &result);
+      return tensorflow::wrap(result);
+    }
+    tensor_handles.push_back(
+        tensorflow::TensorHandleFromInterface(unwrapped_handle));
+  }
+  tensorflow::EagerContext* context =
+      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
+  tensorflow::TensorHandle* handle = nullptr;
+  status->status = tensorflow::TensorHandle::CreatePackedHandle(
+      std::move(tensor_handles), context, &handle);
+  return tensorflow::wrap(handle);
+}
+
+void TFE_ContextSetSoftDevicePlacement(TFE_Context* ctx, unsigned char enable,
+                                       TF_Status* status) {
+  tensorflow::unwrap(ctx)->SetAllowSoftPlacement(enable);
+}
+
+void TFE_ContextSetLogDevicePlacement(TFE_Context* ctx, unsigned char enable,
+                                      TF_Status* status) {
+  tensorflow::unwrap(ctx)->SetLogDevicePlacement(enable);
+}
+
+void TFE_ContextSetRunEagerOpAsFunction(TFE_Context* ctx, unsigned char enable,
+                                        TF_Status* status) {
+  tensorflow::unwrap(ctx)->SetRunEagerOpAsFunction(enable);
+}
+
+void TFE_ContextSetJitCompileRewrite(TFE_Context* ctx, unsigned char enable,
+                                     TF_Status* status) {
+  tensorflow::unwrap(ctx)->SetJitCompileRewrite(enable);
+}
+
+const char* TFE_TensorHandleDeviceType(TFE_TensorHandle* h, TF_Status* status) {
+  if (h == nullptr) {
+    status->status = absl::InvalidArgumentError("Invalid handle");
+    return nullptr;
+  }
+  return tensorflow::unwrap(h)->DeviceType(&status->status);
+}
+
+int TFE_TensorHandleDeviceID(TFE_TensorHandle* h, TF_Status* status) {
+  if (h == nullptr) {
+    status->status = absl::InvalidArgumentError("Invalid handle");
+    return -1;
+  }
+  return tensorflow::unwrap(h)->DeviceId(&status->status);
+}
+
+TF_CAPI_EXPORT extern void TFE_TensorHandleGetStatus(TFE_TensorHandle* h,
+                                                     TF_Status* status) {
+  status->status = tensorflow::unwrap(h)->TensorHandleStatus();
+}
+
+void TFE_GetExecutedOpNames(TFE_Context* ctx, TF_Buffer* buf,
+                            TF_Status* status) {
+  const std::vector<std::string>& op_names =
+      tensorflow::unwrap(ctx)->GetLoggedOpsTestonly();
+
+  std::ostringstream op_names_oss;
+  for (const auto& op : op_names) {
+    op_names_oss << op << ", ";
+  }
+  const std::string& op_names_str = op_names_oss.str();
+  void* data = tensorflow::port::Malloc(op_names_str.length());
+  op_names_str.copy(static_cast<char*>(data), op_names_str.length(), 0);
+  buf->data = data;
+  buf->length = op_names_str.length();
+  buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+  status->status = absl::OkStatus();
+}
+
+void TFE_SetLogicalCpuDevices(TFE_Context* ctx, int num_cpus,
+                              const char* prefix, TF_Status* status) {
+  std::vector<std::unique_ptr<tensorflow::Device>> devices;
+
+  if (prefix == nullptr || strlen(prefix) == 0)
+    prefix = "/job:localhost/replica:0/task:0";
+
+  tensorflow::SessionOptions sess_options;
+  (*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 {
+      ++d;
+    }
+  }
+
+  status->status = tensorflow::unwrap(ctx)->AddDevices(std::move(devices));
+}
+
+void TFE_InsertConfigKeyValue(TFE_Context* ctx, const char* key,
+                              const char* value, TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status = absl::FailedPreconditionError(
+        "Coordination service agent is not enabled.");
+    return;
+  }
+  status->status = coord_agent->InsertKeyValue(key, value);
+}
+
+void TFE_GetConfigKeyValue(TFE_Context* ctx, const char* key,
+                           int64_t timeout_in_ms, TF_Buffer* value_buf,
+                           TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status =
+        absl::FailedPreconditionError("Coordination service is not enabled.");
+    return;
+  }
+  absl::Duration timeout;
+  if (timeout_in_ms > 0) {
+    timeout = absl::Milliseconds(timeout_in_ms);
+  } else {
+    // Block until the key-value is set or the worker shuts down.
+    timeout = absl::InfiniteDuration();
+  }
+  auto status_or_value = coord_agent->GetKeyValue(key, timeout);
+  status->status = status_or_value.status();
+  if (!status_or_value.ok()) return;
+
+  const std::string& value_string = status_or_value.value();
+  void* data = tensorflow::port::Malloc(value_string.length());
+  value_string.copy(static_cast<char*>(data), value_string.length(), 0);
+  value_buf->data = data;
+  value_buf->length = value_string.length();
+  value_buf->data_deallocator = [](void* data, size_t length) {
+    tensorflow::port::Free(data);
+  };
+}
+
+void TFE_DeleteConfigKeyValue(TFE_Context* ctx, const char* key,
+                              TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status =
+        absl::FailedPreconditionError("Coordination service is not enabled.");
+    return;
+  }
+  status->status = coord_agent->DeleteKeyValue(key);
+}
+
+void TFE_ReportErrorToCluster(TFE_Context* ctx, int error_code,
+                              const char* error_message, TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status =
+        absl::FailedPreconditionError("Coordination service is not enabled.");
+    return;
+  }
+  absl::Status s(static_cast<absl::StatusCode>(error_code), error_message);
+  status->status = coord_agent->ReportError(s);
+}
+
+void TFE_GetTaskStates(TFE_Context* ctx, const TF_Buffer& tasks, void* states,
+                       TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status =
+        absl::FailedPreconditionError("Coordination service is not enabled.");
+    return;
+  }
+  std::vector<tensorflow::CoordinatedTask> task_vec(tasks.length);
+  auto* task_iter = static_cast<const tensorflow::CoordinatedTask*>(tasks.data);
+  for (size_t i = 0; i < tasks.length; ++i) {
+    task_vec[i].set_job_name(task_iter->job_name());
+    task_vec[i].set_task_id(task_iter->task_id());
+    ++task_iter;
+  }
+  auto results = coord_agent->GetTaskState(task_vec);
+  if (!results.ok()) {
+    status->status = results.status();
+    return;
+  }
+  auto* state_iter = static_cast<TF_Status*>(states);
+  for (size_t i = 0; i < tasks.length; ++i) {
+    const auto& result = (*results)[i];
+    TF_Status s;
+    TF_SetStatus(&s, static_cast<TF_Code>(result.error_code()),
+                 std::string(result.error_message()).data());
+    if (TF_GetCode(&s) != TF_Code::TF_OK) {
+      tensorflow::CoordinationServiceError error;
+      *error.mutable_source_task() = result.error_payload().source_task();
+      TF_SetPayload(&s, tensorflow::CoordinationErrorPayloadKey().data(),
+                    error.SerializeAsString().c_str());
+    }
+    *state_iter = std::move(s);
+    ++state_iter;
+  }
+  status->status = absl::OkStatus();
+}
+
+void TFE_WaitAtBarrier(TFE_Context* ctx, const char* barrier_id,
+                       int64_t barrier_timeout_in_ms, TF_Status* status) {
+  tensorflow::ImmediateExecutionDistributedManager* dist_mgr =
+      tensorflow::unwrap(ctx)->GetDistributedManager();
+  tsl::CoordinationServiceAgent* coord_agent =
+      dist_mgr->GetCoordinationServiceAgent();
+  if (coord_agent == nullptr) {
+    status->status =
+        absl::FailedPreconditionError("Coordination service is not enabled.");
+    return;
+  }
+  status->status = coord_agent->WaitAtBarrier(
+      barrier_id, absl::Milliseconds(barrier_timeout_in_ms), {});
+}
+
+void TFE_InitializeLocalOnlyContext(TFE_Context* ctx, int keep_alive_secs,
+                                    const void* proto, size_t proto_len,
+                                    TF_Status* status) {
+  tensorflow::ServerDef server_def;
+  if (!server_def.ParseFromArray(proto, proto_len)) {
+    status->status = absl::InvalidArgumentError(
+        "Invalid tensorflow.ServerDef protocol buffer");
+    return;
+  }
+  status->status =
+      tensorflow::unwrap(ctx)
+          ->GetDistributedManager()
+          ->InitializeLocalOnlyContext(server_def, keep_alive_secs);
+}

tensorflow/c/eager/c_api_experimental.h

@@ -0,0 +1,797 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_H_
+#define TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_H_
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/eager/c_api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Resets `op_to_reset` with `op_or_function_name` and `raw_device_name`. This
+// is for performance optimization by reusing an exiting unused op rather than
+// creating a new op every time. If `raw_device_name` is `NULL` or empty, it
+// does not set the device name. If it's not `NULL`, then it attempts to parse
+// and set the device name. It's effectively `TFE_OpSetDevice`, but it is faster
+// than separately calling it because if the existing op has the same
+// `raw_device_name`, it skips parsing and just leave as it is.
+TF_CAPI_EXPORT extern void TFE_OpReset(TFE_Op* op_to_reset,
+                                       const char* op_or_function_name,
+                                       const char* raw_device_name,
+                                       TF_Status* status);
+
+// Enables only graph collection in RunMetadata on the functions executed from
+// this context.
+TF_CAPI_EXPORT extern void TFE_ContextEnableGraphCollection(TFE_Context* ctx);
+
+// Disables only graph collection in RunMetadata on the functions executed from
+// this context.
+TF_CAPI_EXPORT extern void TFE_ContextDisableGraphCollection(TFE_Context* ctx);
+
+// TODO(fishx): Move these monitoring APIs into a separate file.
+// -----------------------------------------------------------------------------
+// Monitoring Counter APIs.
+// These APIs de-templated monitoring Counter for swig.
+
+typedef struct TFE_MonitoringCounterCell TFE_MonitoringCounterCell;
+
+// Atomically increments the value of the cell. The value must be non-negative.
+TF_CAPI_EXPORT extern void TFE_MonitoringCounterCellIncrementBy(
+    TFE_MonitoringCounterCell* cell, int64_t value);
+
+// Retrieves the current value of the cell.
+TF_CAPI_EXPORT extern int64_t TFE_MonitoringCounterCellValue(
+    TFE_MonitoringCounterCell* cell);
+
+// APIs for Counter without label.
+typedef struct TFE_MonitoringCounter0 TFE_MonitoringCounter0;
+// Returns a new Counter metric object. The caller should manage lifetime of
+// the object. Using duplicate metric name will crash the program with fatal
+// error.
+TF_CAPI_EXPORT extern TFE_MonitoringCounter0* TFE_MonitoringNewCounter0(
+    const char* name, TF_Status* status, const char* description);
+// Deletes the Counter object.
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteCounter0(
+    TFE_MonitoringCounter0* counter);
+// Retrieves the cell from the Counter object. The Counter object will manage
+// lifetime of the cell.
+TF_CAPI_EXPORT extern TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter0(
+    TFE_MonitoringCounter0* counter);
+
+// APIs for Counter with 1 label.
+typedef struct TFE_MonitoringCounter1 TFE_MonitoringCounter1;
+TF_CAPI_EXPORT extern TFE_MonitoringCounter1* TFE_MonitoringNewCounter1(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteCounter1(
+    TFE_MonitoringCounter1* counter);
+TF_CAPI_EXPORT extern TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter1(
+    TFE_MonitoringCounter1* counter, const char* label1);
+
+// APIs for Counter with 2 labels.
+typedef struct TFE_MonitoringCounter2 TFE_MonitoringCounter2;
+TF_CAPI_EXPORT extern TFE_MonitoringCounter2* TFE_MonitoringNewCounter2(
+    const char* name, TF_Status* status, const char* description,
+    const char* label1, const char* label2);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteCounter2(
+    TFE_MonitoringCounter2* counter);
+TF_CAPI_EXPORT extern TFE_MonitoringCounterCell* TFE_MonitoringGetCellCounter2(
+    TFE_MonitoringCounter2* counter, const char* label1, const char* label2);
+
+// -----------------------------------------------------------------------------
+// Monitoring Gauge APIs.
+// These APIs de-templated monitoring Gauge for swig.
+
+typedef struct TFE_MonitoringIntGaugeCell TFE_MonitoringIntGaugeCell;
+
+// Atomically set the value of the cell.
+TF_CAPI_EXPORT extern void TFE_MonitoringIntGaugeCellSet(
+    TFE_MonitoringIntGaugeCell* cell, int64_t value);
+
+// Retrieves the current value of the cell.
+TF_CAPI_EXPORT extern int64_t TFE_MonitoringIntGaugeCellValue(
+    TFE_MonitoringIntGaugeCell* cell);
+
+// APIs for Int Gauge without label.
+typedef struct TFE_MonitoringIntGauge0 TFE_MonitoringIntGauge0;
+TF_CAPI_EXPORT extern TFE_MonitoringIntGauge0* TFE_MonitoringNewIntGauge0(
+    const char* name, TF_Status* out_status, const char* description);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteIntGauge0(
+    TFE_MonitoringIntGauge0* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringIntGaugeCell*
+TFE_MonitoringGetCellIntGauge0(TFE_MonitoringIntGauge0* gauge);
+
+// APIs for Int Gauge with 1 label.
+typedef struct TFE_MonitoringIntGauge1 TFE_MonitoringIntGauge1;
+TF_CAPI_EXPORT extern TFE_MonitoringIntGauge1* TFE_MonitoringNewIntGauge1(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteIntGauge1(
+    TFE_MonitoringIntGauge1* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringIntGaugeCell*
+TFE_MonitoringGetCellIntGauge1(TFE_MonitoringIntGauge1* gauge,
+                               const char* label1);
+
+// APIs for Int Gauge with 2 label.
+typedef struct TFE_MonitoringIntGauge2 TFE_MonitoringIntGauge2;
+TF_CAPI_EXPORT extern TFE_MonitoringIntGauge2* TFE_MonitoringNewIntGauge2(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1, const char* label2);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteIntGauge2(
+    TFE_MonitoringIntGauge2* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringIntGaugeCell*
+TFE_MonitoringGetCellIntGauge2(TFE_MonitoringIntGauge2* gauge,
+                               const char* label1, const char* label2);
+
+typedef struct TFE_MonitoringStringGaugeCell TFE_MonitoringStringGaugeCell;
+TF_CAPI_EXPORT extern void TFE_MonitoringStringGaugeCellSet(
+    TFE_MonitoringStringGaugeCell* cell, const char* value);
+// Retrieves the string value and saves it in the buffer.
+TF_CAPI_EXPORT extern const void TFE_MonitoringStringGaugeCellValue(
+    TFE_MonitoringStringGaugeCell* cell, TF_Buffer* buf);
+
+// APIs for String Gauge without label.
+typedef struct TFE_MonitoringStringGauge0 TFE_MonitoringStringGauge0;
+TF_CAPI_EXPORT extern TFE_MonitoringStringGauge0* TFE_MonitoringNewStringGauge0(
+    const char* name, TF_Status* out_status, const char* description);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteStringGauge0(
+    TFE_MonitoringStringGauge0* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringStringGaugeCell*
+TFE_MonitoringGetCellStringGauge0(TFE_MonitoringStringGauge0* gauge);
+
+// APIs for String Gauge with 1 label.
+typedef struct TFE_MonitoringStringGauge1 TFE_MonitoringStringGauge1;
+TF_CAPI_EXPORT extern TFE_MonitoringStringGauge1* TFE_MonitoringNewStringGauge1(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteStringGauge1(
+    TFE_MonitoringStringGauge1* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringStringGaugeCell*
+TFE_MonitoringGetCellStringGauge1(TFE_MonitoringStringGauge1* gauge,
+                                  const char* label1);
+
+// APIs for String Gauge with 2 label.
+typedef struct TFE_MonitoringStringGauge2 TFE_MonitoringStringGauge2;
+TF_CAPI_EXPORT extern TFE_MonitoringStringGauge2* TFE_MonitoringNewStringGauge2(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1, const char* label2);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteStringGauge2(
+    TFE_MonitoringStringGauge2* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringStringGaugeCell*
+TFE_MonitoringGetCellStringGauge2(TFE_MonitoringStringGauge2* gauge,
+                                  const char* label1, const char* label2);
+
+// APIs for String Gauge with 3 labels.
+typedef struct TFE_MonitoringStringGauge3 TFE_MonitoringStringGauge3;
+TF_CAPI_EXPORT extern TFE_MonitoringStringGauge3* TFE_MonitoringNewStringGauge3(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1, const char* label2, const char* label3);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteStringGauge3(
+    TFE_MonitoringStringGauge3* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringStringGaugeCell*
+TFE_MonitoringGetCellStringGauge3(TFE_MonitoringStringGauge3* gauge,
+                                  const char* label1, const char* label2,
+                                  const char* label3);
+
+// APIs for String Gauge with 4 labels.
+typedef struct TFE_MonitoringStringGauge4 TFE_MonitoringStringGauge4;
+TF_CAPI_EXPORT extern TFE_MonitoringStringGauge4* TFE_MonitoringNewStringGauge4(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1, const char* label2, const char* label3,
+    const char* label4);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteStringGauge4(
+    TFE_MonitoringStringGauge4* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringStringGaugeCell*
+TFE_MonitoringGetCellStringGauge4(TFE_MonitoringStringGauge4* gauge,
+                                  const char* label1, const char* label2,
+                                  const char* label3, const char* label4);
+
+typedef struct TFE_MonitoringBoolGaugeCell TFE_MonitoringBoolGaugeCell;
+TF_CAPI_EXPORT extern void TFE_MonitoringBoolGaugeCellSet(
+    TFE_MonitoringBoolGaugeCell* cell, bool value);
+TF_CAPI_EXPORT extern bool TFE_MonitoringBoolGaugeCellValue(
+    TFE_MonitoringBoolGaugeCell* cell);
+
+// APIs for Bool Gauge without label.
+typedef struct TFE_MonitoringBoolGauge0 TFE_MonitoringBoolGauge0;
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGauge0* TFE_MonitoringNewBoolGauge0(
+    const char* name, TF_Status* out_status, const char* description);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteBoolGauge0(
+    TFE_MonitoringBoolGauge0* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGaugeCell*
+TFE_MonitoringGetCellBoolGauge0(TFE_MonitoringBoolGauge0* gauge);
+
+// APIs for Bool Gauge with 1 label.
+typedef struct TFE_MonitoringBoolGauge1 TFE_MonitoringBoolGauge1;
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGauge1* TFE_MonitoringNewBoolGauge1(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteBoolGauge1(
+    TFE_MonitoringBoolGauge1* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGaugeCell*
+TFE_MonitoringGetCellBoolGauge1(TFE_MonitoringBoolGauge1* gauge,
+                                const char* label1);
+
+// APIs for Bool Gauge with 2 label.
+typedef struct TFE_MonitoringBoolGauge2 TFE_MonitoringBoolGauge2;
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGauge2* TFE_MonitoringNewBoolGauge2(
+    const char* name, TF_Status* out_status, const char* description,
+    const char* label1, const char* label2);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteBoolGauge2(
+    TFE_MonitoringBoolGauge2* gauge);
+TF_CAPI_EXPORT extern TFE_MonitoringBoolGaugeCell*
+TFE_MonitoringGetCellBoolGauge2(TFE_MonitoringBoolGauge2* gauge,
+                                const char* label1, const char* label2);
+
+// -----------------------------------------------------------------------------
+// Monitoring Sampler APIs.
+// These APIs de-templated monitoring Sampler for swig.
+
+typedef struct TFE_MonitoringSamplerCell TFE_MonitoringSamplerCell;
+
+// Atomically add the value of the cell.
+TF_CAPI_EXPORT extern void TFE_MonitoringSamplerCellAdd(
+    TFE_MonitoringSamplerCell* cell, double value);
+
+// Retrieves the current value of the cell. The return value is a HistogramProto
+// saved in the buffer.
+TF_CAPI_EXPORT extern void TFE_MonitoringSamplerCellValue(
+    TFE_MonitoringSamplerCell* cell, TF_Buffer* buf);
+
+// APIs for sampler buckets
+typedef struct TFE_MonitoringBuckets TFE_MonitoringBuckets;
+TF_CAPI_EXPORT extern TFE_MonitoringBuckets*
+TFE_MonitoringNewExponentialBuckets(double scale, double growth_factor,
+                                    int bucket_count);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteBuckets(
+    TFE_MonitoringBuckets* buckets);
+
+// APIs for Sampler without label.
+typedef struct TFE_MonitoringSampler0 TFE_MonitoringSampler0;
+TF_CAPI_EXPORT extern TFE_MonitoringSampler0* TFE_MonitoringNewSampler0(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* out_status,
+    const char* description);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteSampler0(
+    TFE_MonitoringSampler0* sampler);
+TF_CAPI_EXPORT extern TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler0(
+    TFE_MonitoringSampler0* sampler);
+
+// APIs for Sampler with 1 label.
+typedef struct TFE_MonitoringSampler1 TFE_MonitoringSampler1;
+TF_CAPI_EXPORT extern TFE_MonitoringSampler1* TFE_MonitoringNewSampler1(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* out_status,
+    const char* description, const char* label1);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteSampler1(
+    TFE_MonitoringSampler1* sampler);
+TF_CAPI_EXPORT extern TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler1(
+    TFE_MonitoringSampler1* sampler, const char* label1);
+
+// APIs for Sampler with 2 label.
+typedef struct TFE_MonitoringSampler2 TFE_MonitoringSampler2;
+TF_CAPI_EXPORT extern TFE_MonitoringSampler2* TFE_MonitoringNewSampler2(
+    const char* name, TFE_MonitoringBuckets* buckets, TF_Status* out_status,
+    const char* description, const char* label1, const char* label2);
+TF_CAPI_EXPORT extern void TFE_MonitoringDeleteSampler2(
+    TFE_MonitoringSampler2* sampler);
+TF_CAPI_EXPORT extern TFE_MonitoringSamplerCell* TFE_MonitoringGetCellSampler2(
+    TFE_MonitoringSampler2* sampler, const char* label1, const char* label2);
+
+// Sets whether to use TFRT
+TF_CAPI_EXPORT extern void TFE_ContextOptionsSetTfrt(TFE_ContextOptions*,
+                                                     bool use_tfrt);
+
+// Returns the context_id from the EagerContext which is used by the
+// EagerService to maintain consistency between client and worker. The
+// context_id is initialized with a dummy value and is later set when the worker
+// is initialized (either locally or remotely). The context_id can change during
+// the process lifetime although this should cause the worker to be
+// reinitialized (e.g. cleared caches) as well.
+TF_CAPI_EXPORT extern uint64_t TFE_GetContextId(TFE_Context* ctx);
+
+// -----------------------------------------------------------------------------
+// Cancellation APIs.
+
+typedef struct TFE_CancellationManager TFE_CancellationManager;
+typedef int64_t TFE_CancellationToken;
+typedef struct TFE_CancelCallback {
+  void (*callback)(void* context);
+  void* context;
+} TFE_CancelCallback;
+TF_CAPI_EXPORT extern TFE_CancellationManager* TFE_NewCancellationManager();
+TF_CAPI_EXPORT extern bool TFE_CancellationManagerIsCancelled(
+    TFE_CancellationManager*);
+TF_CAPI_EXPORT extern bool TFE_CancellationManagerIsCancelling(
+    TFE_CancellationManager*);
+TF_CAPI_EXPORT extern void TFE_CancellationManagerStartCancel(
+    TFE_CancellationManager*);
+TF_CAPI_EXPORT extern TFE_CancellationToken TFE_CancellationManagerGetToken(
+    TFE_CancellationManager*);
+TF_CAPI_EXPORT extern bool TFE_CancellationManagerRegisterCallback(
+    TFE_CancellationManager*, TFE_CancellationToken token,
+    const TFE_CancelCallback* c_callback, const char* callback_name);
+TF_CAPI_EXPORT extern bool TFE_CancellationManagerDeregisterCallback(
+    TFE_CancellationManager*, TFE_CancellationToken token);
+TF_CAPI_EXPORT extern bool TFE_CancellationManagerTryDeregisterCallback(
+    TFE_CancellationManager*, TFE_CancellationToken token);
+TF_CAPI_EXPORT extern void TFE_DeleteCancellationManager(
+    TFE_CancellationManager*);
+
+// Associates the given `cancellation_manager` with `op`, so that invoking
+// `TFE_CancellationManagerStartCancel(cancellation_manager)` will cancel the
+// execution of `op`.
+typedef struct TFE_CancellationManager TFE_CancellationManager;
+TF_CAPI_EXPORT extern void TFE_OpSetCancellationManager(
+    TFE_Op* op, TFE_CancellationManager* cancellation_manager,
+    TF_Status* status);
+
+// -----------------------------------------------------------------------------
+// Eager Executor APIs.
+typedef struct TFE_Executor TFE_Executor;
+
+// Creates a new eager Executor. Nodes in one executor are guaranteed to be
+// executed in sequence. Assigning nodes to different executors allows executing
+// nodes in parallel.
+// in_flight_nodes_limit: when is_async is true, this value controls the
+// maximum number of in flight async nodes. Enqueuing of additional async ops
+// after the limit is reached blocks until some inflight nodes finishes.
+// The effect is bounding the memory held by inflight TensorHandles that are
+// referenced by the inflight nodes.
+// A recommended value has not been established.
+// A value of 0 removes the limit, which is the behavior of TensorFlow 2.11.
+// When is_async is false, the value is ignored.
+TF_CAPI_EXPORT extern TFE_Executor* TFE_NewExecutor(
+    bool is_async, bool enable_streaming_enqueue, int in_flight_nodes_limit);
+
+// Deletes the eager Executor without waiting for enqueued nodes. Please call
+// TFE_ExecutorWaitForAllPendingNodes before calling this API if you want to
+// make sure all nodes are finished.
+TF_CAPI_EXPORT extern void TFE_DeleteExecutor(TFE_Executor*);
+
+// Returns true if the executor is in async mode.
+TF_CAPI_EXPORT extern bool TFE_ExecutorIsAsync(TFE_Executor*);
+
+// Causes the calling thread to block till all ops dispatched in this executor
+// have been executed. Note that "execution" here refers to kernel execution /
+// scheduling of copies, etc. Similar to sync execution, it doesn't guarantee
+// that lower level device queues (like GPU streams) have been flushed.
+//
+// This call may not block for execution of ops enqueued concurrently with this
+// call.
+TF_CAPI_EXPORT extern void TFE_ExecutorWaitForAllPendingNodes(
+    TFE_Executor*, TF_Status* status);
+
+// When an error happens, any pending operations are discarded, and newly issued
+// ops return an error. This call clears the error state and re-enables
+// execution of newly issued ops.
+//
+// Note that outputs of discarded ops remain in a corrupt state and should not
+// be used for future calls.
+// TODO(agarwal): mark the affected handles and raise errors if they are used.
+TF_CAPI_EXPORT extern void TFE_ExecutorClearError(TFE_Executor*);
+
+// Sets a custom Executor for the current thread. All nodes created by this
+// thread will be added to this Executor. It will override the current executor.
+TF_CAPI_EXPORT extern void TFE_ContextSetExecutorForThread(TFE_Context*,
+                                                           TFE_Executor*);
+
+// Returns the Executor for the current thread.
+TF_CAPI_EXPORT extern TFE_Executor* TFE_ContextGetExecutorForThread(
+    TFE_Context*);
+
+// -----------------------------------------------------------------------------
+// Dynamic cluster API.
+
+// Update an existing context with a new set of servers defined in a ServerDef
+// proto. Servers can be added to and removed from the list of remote workers
+// in the context. A New set of servers identified by the ServerDef must be up
+// when the context is updated.
+//
+// This API is for experimental usage and may be subject to change.
+TF_CAPI_EXPORT extern void TFE_ContextUpdateServerDef(TFE_Context* ctx,
+                                                      int keep_alive_secs,
+                                                      const void* proto,
+                                                      size_t proto_len,
+                                                      TF_Status* status);
+
+// This API is for experimental usage and may be subject to change.
+TF_CAPI_EXPORT extern void TFE_ContextUpdateServerDefWithTimeout(
+    TFE_Context* ctx, int keep_alive_secs, const void* proto, size_t proto_len,
+    int64_t init_timeout_in_ms, TF_Status* status);
+
+// This API is for experimental usage and may be subject to change.
+TF_CAPI_EXPORT extern void TFE_ContextSetServerDefWithTimeout(
+    TFE_Context* ctx, int keep_alive_secs, const void* proto, size_t proto_len,
+    int64_t init_timeout_in_ms, TF_Status* status,
+    bool clear_existing_contexts);
+
+// Set server def with retries and timeout. This is helpful for fault-tolerant
+// initial connection in high-preemption environments, such as
+// ParameterServerStrategy training.
+// This API is for experimental usage and may be subject to change.
+TF_CAPI_EXPORT extern void TFE_ContextSetServerDefWithTimeoutAndRetries(
+    TFE_Context* ctx, int keep_alive_secs, const void* proto, size_t proto_len,
+    int64_t init_timeout_in_ms, int retries, TF_Status* status,
+    bool clear_existing_contexts);
+
+// Checks whether a remote worker is alive or not. This will return true even if
+// the context doesn't exist on the remote worker.
+TF_CAPI_EXPORT extern bool TFE_ContextCheckAlive(TFE_Context* ctx,
+                                                 const char* worker_name,
+                                                 TF_Status* status);
+
+// Sync pending nodes in local executors (including the context default executor
+// and thread executors) and streaming requests to remote executors, and get the
+// combined status.
+TF_CAPI_EXPORT extern void TFE_ContextAsyncWait(TFE_Context* ctx,
+                                                TF_Status* status);
+
+// This function will block till the operation that produces `h` has
+// completed. This is only valid on local TFE_TensorHandles. The pointer
+// returned will be on the device in which the TFE_TensorHandle resides (so e.g.
+// for a GPU tensor this will return a pointer to GPU memory). The pointer is
+// only guaranteed to be valid until TFE_DeleteTensorHandle is called on this
+// TensorHandle. Only supports POD data types.
+TF_CAPI_EXPORT extern void* TFE_TensorHandleDevicePointer(TFE_TensorHandle*,
+                                                          TF_Status*);
+
+// This function will block till the operation that produces `h` has
+// completed. This is only valid on local TFE_TensorHandles. Returns the size in
+// bytes of the memory pointed to by the device pointer returned above.
+TF_CAPI_EXPORT extern size_t TFE_TensorHandleDeviceMemorySize(TFE_TensorHandle*,
+                                                              TF_Status*);
+
+// Creates a new TensorHandle from memory residing in the physical device
+// device_name. Takes ownership of the memory, and will call deleter to release
+// it after TF no longer needs it or in case of error.
+//
+// Custom devices must use TFE_NewCustomDeviceTensorHandle instead.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandleFromDeviceMemory(
+    TFE_Context* ctx, const char* device_name, TF_DataType, const int64_t* dims,
+    int num_dims, void* data, size_t len,
+    void (*deallocator)(void* data, size_t len, void* arg),
+    void* deallocator_arg, TF_Status* status);
+
+// Retrieves the address space (i.e. job, replia, task) of the local host and
+// saves it in the buffer.
+TF_CAPI_EXPORT extern void TFE_HostAddressSpace(TFE_Context* ctx,
+                                                TF_Buffer* buf);
+
+// APIs for generically dealing with op attributes (e.g. when forwarding them
+// through custom device implementations).
+//
+// TODO(allenl): Currently these are black boxes, but we should have some way to
+// inspect values. This would let people e.g. copy over most attributes and then
+// modify some based on their values.
+
+// A reference to an op's name -> attribute mapping
+typedef struct TFE_OpAttrs TFE_OpAttrs;
+
+// Fetch a reference to `op`'s attributes. The returned reference is only valid
+// while `op` is alive.
+TF_CAPI_EXPORT extern const TFE_OpAttrs* TFE_OpGetAttrs(const TFE_Op* op);
+// Add attributes in `attrs` to `op`.
+//
+// Does not overwrite or update existing attributes, but adds new ones.
+TF_CAPI_EXPORT extern void TFE_OpAddAttrs(TFE_Op* op, const TFE_OpAttrs* attrs);
+
+// Serialize `attrs` as a tensorflow::NameAttrList protocol buffer (into `buf`),
+// containing the op name and a map of its attributes.
+TF_CAPI_EXPORT extern void TFE_OpAttrsSerialize(const TFE_OpAttrs* attrs,
+                                                TF_Buffer* buf,
+                                                TF_Status* status);
+
+// Set an op's attribute from a serialized AttrValue protocol buffer.
+//
+// Analogous to TF_SetAttrValueProto for building graph operations.
+TF_CAPI_EXPORT extern void TFE_OpSetAttrValueProto(const TFE_Op* op,
+                                                   const char* attr_name,
+                                                   const void* proto,
+                                                   size_t proto_len,
+                                                   TF_Status* status);
+
+// TODO(b/166642410): It would be nice, for custom devices and for other users,
+// to have a non-string representation of devices (TF_Device) extracted from
+// tensors/ops/etc. and usable in APIs like OpSetDevice/ResetOp/etc.
+
+#define TFE_CUSTOM_DEVICE_VERSION 4
+
+// Struct to be filled in. Functions are required except where indicated.
+typedef struct TFE_CustomDevice {
+  int version = TFE_CUSTOM_DEVICE_VERSION;
+  // Method to copy a tensor to the custom device.
+  TFE_TensorHandle* (*copy_tensor_to_device)(TFE_Context* context,
+                                             TFE_TensorHandle* tensor,
+                                             TF_Status* status,
+                                             void* device_info);
+
+  // Method to copy a tensor from the custom device to a target device.
+  TFE_TensorHandle* (*copy_tensor_from_device)(TFE_Context* context,
+                                               TFE_TensorHandle* tensor,
+                                               const char* target_device_name,
+                                               TF_Status* status,
+                                               void* device_info);
+
+  // Method to execute an operation.
+  //
+  // Arguments provide enough information to reconstruct the original `TFE_Op`,
+  // or construct a transformed version, by inspecting the passed `op`.
+  //
+  // TFE_OpGetDevice(op) records the original placement of the operation. It may
+  // be an empty string if no device was explicitly requested, but will
+  // otherwise be the name of this custom device. Ops are placed onto a custom
+  // device if any of their inputs are on that custom device, but custom devices
+  // are free to set a bad status in order to require explicit placement.
+  void (*execute)(const TFE_Op* op, int* num_outputs,
+                  TFE_TensorHandle** outputs, TF_Status* s, void* device_info);
+
+  // Method to delete a device.
+  void (*delete_device)(void* device_info);
+
+  // Implements TFE_CreatePackedTensorHandle when one of `handles` is on this
+  // custom device.
+  //
+  // Many devices will want to simply return an "unimplemented" status
+  // here. This is the default behavior if `pack` is null when passed to
+  // TFE_RegisterCustomDevice.
+  TFE_TensorHandle* (*pack)(TFE_Context* context, TFE_TensorHandle** handles,
+                            int num_handles, TF_Status* s,
+                            void* device_info) = nullptr;
+
+  // Pins the op to `device` based on inputs to `op`. Returns true
+  // signifying to pin to the current custom device. Returns false
+  // to pin to the physical device.
+  //
+  // This function is guaranteed to be called only when all of the custom-device
+  // inputs are on this device.
+  bool (*shall_pin_to_this_device)(const TFE_Op* op, TF_Status* s) = nullptr;
+} TFE_CustomDevice;
+
+// Registers a custom device for use with eager execution.
+//
+// Eager operations may be placed on this device, e.g.  `with
+// tf.device("CUSTOM"):` from Python if `device_name` for this call is
+// "/job:localhost/replica:0/task:0/device:CUSTOM:0".
+//
+// The custom device defines copy operations for moving TensorHandles on and
+// off, and an execution operation for named operations. Often execution will
+// simply wrap op execution on one or more physical devices.
+//
+// device_info is an opaque caller-defined type stored with the custom device
+// which is passed to the functions referenced in the TFE_CustomDevice struct
+// `device` (execute, delete_device, etc.). It can for example contain the
+// names of wrapped devices.
+//
+// There are currently no graph semantics implemented for registered custom
+// devices, so executing tf.functions which contain operations placed on the
+// custom devices will fail.
+//
+// `device_name` must not name an existing physical or custom device. It must
+// follow the format:
+//
+//    /job:<name>/replica:<replica>/task:<task>/device:<type>:<device_num>
+//
+// If the device is successfully registered, `status` is set to TF_OK. Otherwise
+// the device is not usable. In case of a bad status, `device.delete_device` is
+// still called on `device_info` (i.e. the caller does not retain ownership).
+//
+// This API is highly experimental, and in particular is expected to change when
+// it starts supporting operations with attributes and when tf.function support
+// is added.
+TF_CAPI_EXPORT extern void TFE_RegisterCustomDevice(TFE_Context* ctx,
+                                                    TFE_CustomDevice device,
+                                                    const char* device_name,
+                                                    void* device_info,
+                                                    TF_Status* status);
+
+// Returns whether `device_name` maps to a registered custom device.
+TF_CAPI_EXPORT extern bool TFE_IsCustomDevice(TFE_Context* ctx,
+                                              const char* device_name);
+
+// Struct to be filled in to define a custom device tensor handle. Fields are
+// required except where indicated.
+typedef struct TFE_CustomDeviceTensorHandleMethods {
+  int version = TFE_CUSTOM_DEVICE_VERSION;
+
+  // Computes the rank of the tensor handle.
+  //
+  // Shapes are specified via callbacks because retrieving the shape of a tensor
+  // is a blocking operation for async eager; custom devices should avoid
+  // retrieving shapes of tensors they wrap until the custom device tensor's
+  // shape is explicitly requested where possible.
+  int (*num_dims)(void* data, TF_Status* status);
+
+  // Computes the axis length at `dim_index`.
+  int64_t (*dim)(void* data, int dim_index, TF_Status* status);
+
+  void (*deallocator)(void* data);
+
+  // Summarizes the value of this tensor. The caller takes ownership of the
+  // returned buffer. If `status` is not TF_OK, instead returns a null pointer.
+  //
+  // Does not include the shape and dtype of the tensor (which is generally
+  // appended later), but should include any information specific to this custom
+  // device which would be useful for debugging.
+  //
+  // Optional. If null, defaults to resolving the TFE_TensorHandle into a
+  // TF_Tensor and summarizing that.
+  TF_Buffer* (*summarize)(void* data, TF_Status* status) = nullptr;
+} TFE_CustomDeviceTensorHandle;
+
+// Creates a new TensorHandle from memory residing in a custom device. Takes
+// ownership of the memory pointed to by `tensor_handle_data`, and calls
+// `methods.deallocator` to release it after TF no longer needs it or in case of
+// an error.
+//
+// This call is similar to `TFE_NewTensorHandleFromDeviceMemory`, but supports
+// custom devices instead of physical devices and does not require blocking
+// waiting for exact shapes.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewCustomDeviceTensorHandle(
+    TFE_Context*, const char* device_name, TF_DataType, void* data,
+    TFE_CustomDeviceTensorHandle methods, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_ContextGetFunctionDef(TFE_Context* ctx,
+                                                     const char* function_name,
+                                                     TF_Buffer* buf,
+                                                     TF_Status* status);
+
+// Get GraphDebugInfo containing stack traces mapping to node names
+TF_CAPI_EXPORT extern void TFE_ContextGetGraphDebugInfo(
+    TFE_Context* ctx, const char* function_name, TF_Buffer* buf,
+    TF_Status* status);
+
+// Extracts a TF_Function from the context.
+// Must call TF_DeleteFunction on the returned value.
+TF_CAPI_EXPORT extern TF_Function* TFE_ContextGetFunction(TFE_Context* ctx,
+                                                          const char* name,
+                                                          TF_Status* status);
+
+// Allocate and return a new Tensor on the host.
+//
+// The caller must set the Tensor values by writing them to the pointer returned
+// by TF_TensorData with length TF_TensorByteSize.
+TF_CAPI_EXPORT extern TF_Tensor* TFE_AllocateHostTensor(TFE_Context* ctx,
+                                                        TF_DataType dtype,
+                                                        const int64_t* dims,
+                                                        int num_dims,
+                                                        TF_Status* status);
+
+// Given a Tensor, wrap it with a TensorHandle
+//
+// Similar to TFE_NewTensorHandle, but includes a pointer to the TFE_Context.
+// The context should be identical to that of the Tensor.
+TF_CAPI_EXPORT TFE_TensorHandle* TFE_NewTensorHandleFromTensor(
+    TFE_Context* ctx, TF_Tensor* t, TF_Status* status);
+
+// Create a packed TensorHandle with the given list of TensorHandles.
+// If `handles` are on the same device, assign the same device to the packed
+// handle; if `handles` are on different deivces, assign a CompositeDevice to
+// it.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_CreatePackedTensorHandle(
+    TFE_Context* ctx, TFE_TensorHandle** handles, int* num_handles,
+    TF_Status* status);
+
+// Configure soft device placement policy for the eager executor. Note this
+// policy is applied to any subsequent op executions.
+TF_CAPI_EXPORT void TFE_ContextSetSoftDevicePlacement(TFE_Context* ctx,
+                                                      unsigned char enable,
+                                                      TF_Status* status);
+
+// Configure device placement policy logging for the eager executor. Note this
+// policy is applied to any subsequent op executions.
+TF_CAPI_EXPORT void TFE_ContextSetLogDevicePlacement(TFE_Context* ctx,
+                                                     unsigned char enable,
+                                                     TF_Status* status);
+
+// Enables running eager ops as function.
+TF_CAPI_EXPORT void TFE_ContextSetRunEagerOpAsFunction(TFE_Context* ctx,
+                                                       unsigned char enable,
+                                                       TF_Status* status);
+
+// Enables rewrite jit_compile functions.
+TF_CAPI_EXPORT void TFE_ContextSetJitCompileRewrite(TFE_Context* ctx,
+                                                    unsigned char enable,
+                                                    TF_Status* status);
+
+// Returns the device type of the operation that produced `h`.
+TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceType(
+    TFE_TensorHandle* h, TF_Status* status);
+
+// Returns the device ID of the operation that produced `h`.
+TF_CAPI_EXPORT extern int TFE_TensorHandleDeviceID(TFE_TensorHandle* h,
+                                                   TF_Status* status);
+
+// Returns the status for the tensor handle. In TFRT, a tensor handle can carry
+// error info if error happens. If so, the status will be set with the error
+// info. If not, status will be set as OK.
+TF_CAPI_EXPORT extern void TFE_TensorHandleGetStatus(TFE_TensorHandle* h,
+                                                     TF_Status* status);
+
+// Get a comma-separated list of op names executed in graph functions dispatched
+// to `ctx`. This feature is currently only enabled for TFRT debug builds, for
+// performance and simplicity reasons.
+TF_CAPI_EXPORT extern void TFE_GetExecutedOpNames(TFE_Context* ctx,
+                                                  TF_Buffer* buf,
+                                                  TF_Status* status);
+
+// Set logical devices to the context's device manager.
+// If logical devices are already configured at context initialization
+// through TFE_ContextOptions, this method should not be called.
+TF_CAPI_EXPORT extern void TFE_SetLogicalCpuDevices(TFE_Context* ctx,
+                                                    int num_cpus,
+                                                    const char* prefix,
+                                                    TF_Status* status);
+
+// Set configuration key and value using coordination service.
+// If coordination service is enabled, the key-value will be stored on the
+// leader and become accessible to all workers in the cluster.
+// Currently, a config key can only be set with one value, and subsequently
+// setting the same key will lead to errors.
+//
+// Note that the key-values are only expected to be used for cluster
+// configuration data, and should not be used for storing a large amount of data
+// or being accessed very frequently.
+TF_CAPI_EXPORT extern void TFE_InsertConfigKeyValue(TFE_Context* ctx,
+                                                    const char* key,
+                                                    const char* value,
+                                                    TF_Status* status);
+
+// Get configuration key and value using coordination service.
+// The config key must be set before getting its value. Getting value of
+// non-existing config keys will result in errors.
+// If `timeout_in_ms=0`, this call will block until the key-value is set or the
+// worker shuts down.
+TF_CAPI_EXPORT extern void TFE_GetConfigKeyValue(TFE_Context* ctx,
+                                                 const char* key,
+                                                 int64_t timeout_in_ms,
+                                                 TF_Buffer* value_buf,
+                                                 TF_Status* status);
+
+// Delete configuration key-value. If `key` is a directory, recursively clean up
+// all key-values under the path specified by `key`.
+TF_CAPI_EXPORT extern void TFE_DeleteConfigKeyValue(TFE_Context* ctx,
+                                                    const char* key,
+                                                    TF_Status* status);
+
+// Report error (specified by error_code and error_message) to other tasks in
+// the cluster.
+TF_CAPI_EXPORT extern void TFE_ReportErrorToCluster(TFE_Context* ctx,
+                                                    int error_code,
+                                                    const char* error_message,
+                                                    TF_Status* status);
+
+// Get task states from the Coordination Service.
+TF_CAPI_EXPORT extern void TFE_GetTaskStates(TFE_Context* ctx,
+                                             const TF_Buffer& tasks,
+                                             void* states, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_WaitAtBarrier(TFE_Context* ctx,
+                                             const char* barrier_id,
+                                             int64_t barrier_timeout_in_ms,
+                                             TF_Status* status);
+
+TF_CAPI_EXPORT extern void TFE_InitializeLocalOnlyContext(TFE_Context* ctx,
+                                                          int keep_alive_secs,
+                                                          const void* proto,
+                                                          size_t proto_len,
+                                                          TF_Status* status);
+
+#ifdef __cplusplus
+} /* end extern "C" */
+#endif
+
+#endif  // TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_H_

tensorflow/c/eager/c_api_experimental_reader.cc

@@ -0,0 +1,42 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_experimental_reader.h"
+
+#include "tensorflow/c/eager/tfe_monitoring_reader_internal.h"
+
+template <typename... LabelType>
+int64_t TFE_MonitoringCounterReader::Read(const LabelType&... labels) {
+  return counter->Read(labels...);
+}
+
+TFE_MonitoringCounterReader* TFE_MonitoringNewCounterReader(const char* name) {
+  auto* result = new TFE_MonitoringCounterReader(name);
+
+  return result;
+}
+
+int64_t TFE_MonitoringReadCounter0(TFE_MonitoringCounterReader* cell_reader) {
+  int64_t result = cell_reader->Read();
+
+  return result;
+}
+
+int64_t TFE_MonitoringReadCounter1(TFE_MonitoringCounterReader* cell_reader,
+                                   const char* label) {
+  int64_t result = cell_reader->Read(label);
+
+  return result;
+}

tensorflow/c/eager/c_api_experimental_reader.h

@@ -0,0 +1,60 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_READER_H_
+#define TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_READER_H_
+
+#include "tensorflow/c/eager/c_api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Test only exports of the monitoring Cell Reader API which allows tests to
+// read current values from streamz counters defined in other modules.
+//
+// The code under test will have created streamz counters like this:
+// auto* streamz = tensorflow::monitoring::Counter<1>::New("name",
+// "description", "label");
+// and then incremented that counter for various values of label:
+// streamz->GetCell("label-value")->IncrementBy(1);
+//
+// The test code can then read and test the value of that counter:
+//
+// auto* reader = TFE_MonitoringNewCounterReader("name");
+// test();
+// int64_t value = TFE_MonitoringReadCounter1(reader, "label-value");
+
+// Opaque handle to a reader.
+typedef struct TFE_MonitoringCounterReader TFE_MonitoringCounterReader;
+
+// Returns a handle to be used for reading values from streamz counter. The
+// counter can have been created with any number of labels.
+TF_CAPI_EXPORT extern TFE_MonitoringCounterReader*
+TFE_MonitoringNewCounterReader(const char* name);
+
+// Reads the value of a counter that was created with 0 labels.
+TF_CAPI_EXPORT extern int64_t TFE_MonitoringReadCounter0(
+    TFE_MonitoringCounterReader*);
+
+// Reads the value of specific cell of a counter that was created with 1 label.
+TF_CAPI_EXPORT extern int64_t TFE_MonitoringReadCounter1(
+    TFE_MonitoringCounterReader*, const char* label_value);
+
+#ifdef __cplusplus
+} /* end extern "C" */
+#endif
+
+#endif  // TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_READER_H_

tensorflow/c/eager/c_api_experimental_reader_test.cc

@@ -0,0 +1,86 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_experimental_reader.h"
+
+#include <cstdint>
+
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TFE_MonitoringCounter0* CreateCounter0(const char* counter_name);
+TFE_MonitoringCounter1* CreateCounter1(const char* counter_name,
+                                       const char* label);
+void IncrementCounter0(TFE_MonitoringCounter0* counter, int64_t delta = 1);
+void IncrementCounter1(TFE_MonitoringCounter1* counter, const char* label,
+                       int64_t delta = 1);
+
+TEST(CAPI, MonitoringCellReader0) {
+  auto counter_name = "test/counter0";
+  auto* counter = CreateCounter0(counter_name);
+  auto* reader = TFE_MonitoringNewCounterReader(counter_name);
+  IncrementCounter0(counter);
+
+  int64_t actual = TFE_MonitoringReadCounter0(reader);
+
+  CHECK_EQ(actual, 1);
+}
+
+TEST(CAPI, MonitoringCellReader1) {
+  auto counter_name = "test/counter1";
+  auto label_name = "test/label";
+  auto* counter = CreateCounter1(counter_name, label_name);
+  auto* reader = TFE_MonitoringNewCounterReader(counter_name);
+  IncrementCounter1(counter, label_name);
+
+  int64_t actual = TFE_MonitoringReadCounter1(reader, label_name);
+
+  CHECK_EQ(actual, 1);
+}
+
+TFE_MonitoringCounter0* CreateCounter0(const char* counter_name) {
+  TF_Status* status = TF_NewStatus();
+  auto* counter =
+      TFE_MonitoringNewCounter0(counter_name, status, "description");
+  TF_DeleteStatus(status);
+  return counter;
+}
+
+void IncrementCounter0(TFE_MonitoringCounter0* counter, int64_t delta) {
+  auto* cell = TFE_MonitoringGetCellCounter0(counter);
+  TFE_MonitoringCounterCellIncrementBy(cell, delta);
+}
+
+TFE_MonitoringCounter1* CreateCounter1(const char* counter_name,
+                                       const char* label) {
+  TF_Status* status = TF_NewStatus();
+  auto* counter =
+      TFE_MonitoringNewCounter1(counter_name, status, "description", label);
+  TF_DeleteStatus(status);
+  return counter;
+}
+
+void IncrementCounter1(TFE_MonitoringCounter1* counter, const char* label,
+                       int64_t delta) {
+  auto* cell = TFE_MonitoringGetCellCounter1(counter, label);
+  TFE_MonitoringCounterCellIncrementBy(cell, delta);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/eager/c_api_experimental_test.cc

@@ -0,0 +1,762 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_experimental.h"
+
+#include <string.h>
+
+#include <string>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/distributed_runtime/server_lib.h"
+#include "tensorflow/core/lib/monitoring/collection_registry.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+
+using tensorflow::string;
+
+namespace tensorflow {
+namespace {
+
+static bool HasSubstr(absl::string_view base, absl::string_view substr) {
+  bool ok = absl::StrContains(base, substr);
+  EXPECT_TRUE(ok) << base << ", expected substring " << substr;
+  return ok;
+}
+
+TEST(CAPI, MonitoringCounter0) {
+  TF_Status* status = TF_NewStatus();
+  auto* counter =
+      TFE_MonitoringNewCounter0("test/counter", status, "description");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  auto* cell = TFE_MonitoringGetCellCounter0(counter);
+  TFE_MonitoringCounterCellIncrementBy(cell, 1);
+  EXPECT_EQ(TFE_MonitoringCounterCellValue(cell), 1);
+  auto* collection_registry = monitoring::CollectionRegistry::Default();
+  monitoring::CollectionRegistry::CollectMetricsOptions options;
+  std::unique_ptr<monitoring::CollectedMetrics> metrics =
+      collection_registry->CollectMetrics(options);
+
+  EXPECT_EQ("test/counter",
+            metrics->point_set_map.at("test/counter")->metric_name);
+  EXPECT_EQ(
+      1, metrics->point_set_map.at("test/counter")->points.at(0)->int64_value);
+
+  TFE_MonitoringCounterCellIncrementBy(cell, 5);
+  EXPECT_EQ(TFE_MonitoringCounterCellValue(cell), 6);
+  metrics = collection_registry->CollectMetrics(options);
+  EXPECT_EQ(
+      6, metrics->point_set_map.at("test/counter")->points.at(0)->int64_value);
+
+  TFE_MonitoringDeleteCounter0(counter);
+  metrics = collection_registry->CollectMetrics(options);
+  EXPECT_EQ(metrics->point_set_map.end(),
+            metrics->point_set_map.find("test/counter"));
+}
+
+TEST(CAPI, MonitoringCounterMultiple) {
+  TF_Status* status = TF_NewStatus();
+  auto* counter1 = TFE_MonitoringNewCounter1("test/counter1", status,
+                                             "description", "label1");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell1 = TFE_MonitoringGetCellCounter1(counter1, "test");
+  TFE_MonitoringCounterCellIncrementBy(cell1, 1);
+  EXPECT_EQ(TFE_MonitoringCounterCellValue(cell1), 1);
+
+  auto* counter2 = TFE_MonitoringNewCounter2("test/counter2", status,
+                                             "description", "label1", "label2");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  auto* cell2 = TFE_MonitoringGetCellCounter2(counter2, "foo", "bar");
+  TFE_MonitoringCounterCellIncrementBy(cell2, 2);
+  EXPECT_EQ(TFE_MonitoringCounterCellValue(cell2), 2);
+
+  TFE_MonitoringDeleteCounter1(counter1);
+  TFE_MonitoringDeleteCounter2(counter2);
+}
+
+TEST(CAPI, MonitoringGauge0) {
+  TF_Status* status = TF_NewStatus();
+  auto* gauge = TFE_MonitoringNewIntGauge0("test/gauge", status, "test");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell = TFE_MonitoringGetCellIntGauge0(gauge);
+  TFE_MonitoringIntGaugeCellSet(cell, 1);
+  EXPECT_EQ(TFE_MonitoringIntGaugeCellValue(cell), 1);
+  auto* collection_registry = monitoring::CollectionRegistry::Default();
+  monitoring::CollectionRegistry::CollectMetricsOptions options;
+  std::unique_ptr<monitoring::CollectedMetrics> metrics =
+      collection_registry->CollectMetrics(options);
+
+  EXPECT_EQ("test/gauge", metrics->point_set_map.at("test/gauge")->metric_name);
+  EXPECT_EQ(1,
+            metrics->point_set_map.at("test/gauge")->points.at(0)->int64_value);
+
+  TFE_MonitoringIntGaugeCellSet(cell, 5);
+  metrics = collection_registry->CollectMetrics(options);
+  EXPECT_EQ(5,
+            metrics->point_set_map.at("test/gauge")->points.at(0)->int64_value);
+  TFE_MonitoringDeleteIntGauge0(gauge);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, MonitoringMultipleGauge) {
+  TF_Status* status = TF_NewStatus();
+  auto* gauge1 =
+      TFE_MonitoringNewBoolGauge1("test/gauge1", status, "test", "label1");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell1 = TFE_MonitoringGetCellBoolGauge1(gauge1, "foo");
+  TFE_MonitoringBoolGaugeCellSet(cell1, true);
+  EXPECT_TRUE(TFE_MonitoringBoolGaugeCellValue(cell1));
+  TFE_MonitoringDeleteBoolGauge1(gauge1);
+
+  auto* gauge2 = TFE_MonitoringNewStringGauge2("test/gauge2", status, "test",
+                                               "label1", "label2");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell2 = TFE_MonitoringGetCellStringGauge2(gauge2, "foo", "bar");
+  TFE_MonitoringStringGaugeCellSet(cell2, "str");
+  auto* buf = new TF_Buffer;
+  TFE_MonitoringStringGaugeCellValue(cell2, buf);
+  std::string data(static_cast<const char*>(buf->data), buf->length);
+  TF_DeleteBuffer(buf);
+  EXPECT_EQ(data, "str");
+  TFE_MonitoringDeleteStringGauge2(gauge2);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, MonitoringSampler0) {
+  TF_Status* status = TF_NewStatus();
+  auto* buckets = TFE_MonitoringNewExponentialBuckets(1.0, 2.0, 2);
+  auto* sampler =
+      TFE_MonitoringNewSampler0("test/sampler", buckets, status, "test");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell = TFE_MonitoringGetCellSampler0(sampler);
+  TFE_MonitoringSamplerCellAdd(cell, 1.0);
+  auto* collection_registry = monitoring::CollectionRegistry::Default();
+  monitoring::CollectionRegistry::CollectMetricsOptions options;
+  std::unique_ptr<monitoring::CollectedMetrics> metrics =
+      collection_registry->CollectMetrics(options);
+
+  EXPECT_EQ("test/sampler",
+            metrics->point_set_map.at("test/sampler")->metric_name);
+  EXPECT_EQ(1.0, metrics->point_set_map.at("test/sampler")
+                     ->points.at(0)
+                     ->histogram_value.sum());
+
+  TFE_MonitoringSamplerCellAdd(cell, 5.0);
+  metrics = collection_registry->CollectMetrics(options);
+  EXPECT_EQ(6.0, metrics->point_set_map.at("test/sampler")
+                     ->points.at(0)
+                     ->histogram_value.sum());
+  TFE_MonitoringDeleteBuckets(buckets);
+  TFE_MonitoringDeleteSampler0(sampler);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, MonitoringMultipleSampler) {
+  TF_Status* status = TF_NewStatus();
+  auto* buckets = TFE_MonitoringNewExponentialBuckets(1.0, 2.0, 2);
+  auto* sampler1 = TFE_MonitoringNewSampler1("test/sampler1", buckets, status,
+                                             "test", "label1");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell1 = TFE_MonitoringGetCellSampler1(sampler1, "foo");
+  TFE_MonitoringSamplerCellAdd(cell1, 1.0);
+  TFE_MonitoringSamplerCellAdd(cell1, 2.0);
+  TF_Buffer* result1 = TF_NewBuffer();
+  TFE_MonitoringSamplerCellValue(cell1, result1);
+  tensorflow::HistogramProto histogram1;
+  EXPECT_TRUE(histogram1.ParseFromString(
+      {reinterpret_cast<const char*>(result1->data), result1->length}));
+  EXPECT_EQ(histogram1.sum(), 3.0);
+  TF_DeleteBuffer(result1);
+  TFE_MonitoringDeleteSampler1(sampler1);
+
+  auto* sampler2 = TFE_MonitoringNewSampler2("test/sampler2", buckets, status,
+                                             "test", "label1", "label2");
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* cell2 = TFE_MonitoringGetCellSampler2(sampler2, "foo", "bar");
+  TFE_MonitoringSamplerCellAdd(cell2, 2.0);
+  TFE_MonitoringSamplerCellAdd(cell2, 3.0);
+  TF_Buffer* result2 = TF_NewBuffer();
+  TFE_MonitoringSamplerCellValue(cell2, result2);
+  tensorflow::HistogramProto histogram2;
+  EXPECT_TRUE(histogram2.ParseFromString(
+      {reinterpret_cast<const char*>(result2->data), result2->length}));
+  EXPECT_EQ(histogram2.sum(), 5.0);
+  TF_DeleteBuffer(result2);
+  TFE_MonitoringDeleteSampler2(sampler2);
+
+  TFE_MonitoringDeleteBuckets(buckets);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, CancellationManager) {
+  TFE_CancellationManager* c_mgr = TFE_NewCancellationManager();
+  EXPECT_FALSE(TFE_CancellationManagerIsCancelled(c_mgr));
+
+  TFE_CancelCallback callback1;
+  callback1.callback = [](void* context) {
+    ADD_FAILURE() << "Callback1 should be deregistered.";
+  };
+  TFE_CancellationToken token1 = TFE_CancellationManagerGetToken(c_mgr);
+  EXPECT_TRUE(TFE_CancellationManagerRegisterCallback(c_mgr, token1, &callback1,
+                                                      "callback1"));
+
+  TFE_CancelCallback callback2;
+  bool callback2_invoked = false;
+  callback2.context = &callback2_invoked;
+  callback2.callback = [](void* context) {
+    *reinterpret_cast<bool*>(context) = true;
+  };
+  TFE_CancellationToken token2 = TFE_CancellationManagerGetToken(c_mgr);
+  EXPECT_TRUE(TFE_CancellationManagerRegisterCallback(c_mgr, token2, &callback2,
+                                                      "callback2"));
+
+  TFE_CancellationToken token3 = TFE_CancellationManagerGetToken(c_mgr);
+  EXPECT_TRUE(TFE_CancellationManagerRegisterCallback(c_mgr, token3, &callback1,
+                                                      "callback3"));
+
+  EXPECT_TRUE(TFE_CancellationManagerDeregisterCallback(c_mgr, token1));
+  EXPECT_TRUE(TFE_CancellationManagerTryDeregisterCallback(c_mgr, token3));
+
+  TFE_CancellationManagerStartCancel(c_mgr);
+  EXPECT_TRUE(TFE_CancellationManagerIsCancelled(c_mgr));
+  EXPECT_TRUE(callback2_invoked);
+  TFE_DeleteCancellationManager(c_mgr);
+}
+
+TEST(CAPI, ExecutorContextDestructionOrder) {
+  TF_Status* status = TF_NewStatus();
+
+  {
+    TFE_ContextOptions* opts = TFE_NewContextOptions();
+    TFE_Context* ctx = TFE_NewContext(opts, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TFE_DeleteContextOptions(opts);
+    TFE_Executor* executor = TFE_NewExecutor(
+        /*is_async=*/false, /*enable_streaming_enqueue=*/true,
+        /*in_flight_nodes_limit=*/0);
+    TFE_ContextSetExecutorForThread(ctx, executor);
+
+    TFE_DeleteContext(ctx);
+    TFE_DeleteExecutor(executor);
+  }
+
+  {
+    TFE_ContextOptions* opts = TFE_NewContextOptions();
+    TFE_Context* ctx = TFE_NewContext(opts, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TFE_DeleteContextOptions(opts);
+    TFE_Executor* executor = TFE_NewExecutor(
+        /*is_async=*/false, /*enable_streaming_enqueue=*/true,
+        /*in_flight_nodes_limit=*/0);
+    TFE_ContextSetExecutorForThread(ctx, executor);
+
+    TFE_DeleteExecutor(executor);
+    TFE_DeleteContext(ctx);
+  }
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, Function_ident_CPU) {
+  // First create a simple identity function.
+  TF_Graph* function_graph = TF_NewGraph();
+  TF_OperationDescription* arg_descr =
+      TF_NewOperation(function_graph, "Placeholder", "arg");
+  TF_SetAttrType(arg_descr, "dtype", TF_INT32);
+  TF_Status* status = TF_NewStatus();
+  TF_Operation* arg = TF_FinishOperation(arg_descr, status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TF_OperationDescription* id_descr =
+      TF_NewOperation(function_graph, "Identity", "id");
+  TF_SetAttrType(id_descr, "T", TF_INT32);
+  TF_AddInput(id_descr, {arg, 0});
+  TF_Operation* id = TF_FinishOperation(id_descr, status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TF_Output input{arg, 0};
+  TF_Output output{id, 0};
+  TF_Function* fn =
+      TF_GraphToFunction(function_graph, "ident", 0, 1, &id, 1, &input, 1,
+                         &output, nullptr, nullptr, "test", status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TF_DeleteGraph(function_graph);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+  TFE_ContextAddFunction(ctx, fn, status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TF_DeleteFunction(fn);
+
+  for (bool async : {false, true, false}) {
+    TFE_Executor* old_executor = TFE_ContextGetExecutorForThread(ctx);
+    TFE_Executor* executor = TFE_NewExecutor(
+        /*is_async=*/async, /*enable_streaming_enqueue=*/true,
+        /*in_flight_nodes_limit=*/0);
+    TFE_ContextSetExecutorForThread(ctx, executor);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+    TF_Tensor* t = TF_AllocateTensor(TF_INT32, nullptr, 0, 1 * sizeof(int32_t));
+    *reinterpret_cast<int32_t*>(TF_TensorData(t)) = 42;
+    TFE_TensorHandle* h = TFE_NewTensorHandle(t, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TF_DeleteTensor(t);
+
+    TFE_Op* op = TFE_NewOp(ctx, "ident", status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TFE_OpAddInput(op, h, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+
+    std::vector<TFE_TensorHandle*> result;
+    result.push_back(nullptr);
+    int num_retvals = 1;
+    TFE_Execute(op, result.data(), &num_retvals, status);
+    TFE_DeleteOp(op);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    ASSERT_EQ(num_retvals, 1);
+
+    TF_Tensor* r = TFE_TensorHandleResolve(result[0], status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    EXPECT_EQ(*reinterpret_cast<int32_t*>(TF_TensorData(r)), 42);
+    TFE_ContextSetExecutorForThread(ctx, old_executor);
+    TFE_ExecutorWaitForAllPendingNodes(executor, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_DeleteExecutor(executor);
+    TFE_DeleteExecutor(old_executor);
+    TFE_DeleteTensorHandle(h);
+    TF_DeleteTensor(r);
+    TFE_DeleteTensorHandle(result[0]);
+  }
+  TFE_ContextRemoveFunction(ctx, "ident", status);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TFE_DeleteContext(ctx);
+  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+  TF_DeleteStatus(status);
+}
+
+void Executor_MatMul_CPU(bool async) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_Executor* old_executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_Executor* executor = TFE_NewExecutor(
+      /*is_async=*/async, /*enable_streaming_enqueue=*/true,
+      /*in_flight_nodes_limit=*/0);
+  TFE_ContextSetExecutorForThread(ctx, executor);
+
+  TFE_TensorHandle* m = TestMatrixTensorHandle(ctx);
+  TFE_Op* matmul = MatMulOp(ctx, m, m);
+  TFE_TensorHandle* retvals[2] = {nullptr, nullptr};
+  int num_retvals = 2;
+  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(1, num_retvals);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteOp(matmul);
+  TFE_DeleteTensorHandle(m);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteTensorHandle(retvals[0]);
+  TFE_ContextSetExecutorForThread(ctx, old_executor);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+  TFE_DeleteExecutor(old_executor);
+  TFE_DeleteContext(ctx);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  float product[4] = {0};
+  EXPECT_EQ(sizeof(product), TF_TensorByteSize(t));
+  memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+  EXPECT_EQ(7, product[0]);
+  EXPECT_EQ(10, product[1]);
+  EXPECT_EQ(15, product[2]);
+  EXPECT_EQ(22, product[3]);
+  TF_DeleteStatus(status);
+}
+TEST(CAPI, Executor_MatMul_CPU) { Executor_MatMul_CPU(false); }
+TEST(CAPI, Executor_MatMul_CPUAsync) { Executor_MatMul_CPU(true); }
+
+void Deleter(void* data, size_t unused, void* tensor_handle) {
+  TFE_DeleteTensorHandle(static_cast<TFE_TensorHandle*>(tensor_handle));
+}
+
+TEST(CAPI, TensorHandleOnDeviceMemory) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_TensorHandle* m = TestMatrixTensorHandle(ctx);
+  TF_Tensor* m_data = TFE_TensorHandleResolve(m, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  float* m_float = static_cast<float*>(TF_TensorData(m_data));
+  TF_DeviceList* devices = TFE_ContextListDevices(ctx, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  int num_devices = TF_DeviceListCount(devices);
+  for (int d = 0; d < num_devices; ++d) {
+    const char* name = TF_DeviceListName(devices, d, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_TensorHandle* copy = TFE_TensorHandleCopyToDevice(m, ctx, name, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    void* data = TFE_TensorHandleDevicePointer(copy, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    size_t size = TFE_TensorHandleDeviceMemorySize(copy, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    int64_t dims[] = {2, 2};
+    TFE_TensorHandle* copy_aliased = TFE_NewTensorHandleFromDeviceMemory(
+        ctx, name, TF_FLOAT, dims, 2, data, size, &Deleter, copy, status);
+    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 =
+        static_cast<const float*>(TF_TensorData(resolved));
+    EXPECT_EQ(0, memcmp(m_float, resolved_data, 4 * sizeof(float)));
+    TF_DeleteTensor(resolved);
+    TFE_DeleteTensorHandle(copy_aliased);  // Note that this will delete copy.
+    TFE_DeleteTensorHandle(on_host);
+  }
+  TF_DeleteDeviceList(devices);
+  TF_DeleteTensor(m_data);
+  TFE_DeleteTensorHandle(m);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+}
+
+TEST(CAPI, TensorHandleNullptr) {
+  TFE_TensorHandle* h = nullptr;
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+
+  const char* device_type = TFE_TensorHandleDeviceType(h, status.get());
+  ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(status.get()));
+  ASSERT_EQ(device_type, nullptr);
+  ASSERT_EQ("Invalid handle", std::string(TF_Message(status.get())));
+
+  TF_SetStatus(status.get(), TF_OK, "");
+
+  int device_id = TFE_TensorHandleDeviceID(h, status.get());
+  ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(status.get()));
+  ASSERT_EQ(device_id, -1);
+  ASSERT_EQ("Invalid handle", std::string(TF_Message(status.get())));
+}
+
+TEST(CAPI, TensorHandleDevices) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status.get());
+  TFE_DeleteContextOptions(opts);
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+
+  TFE_TensorHandle* hcpu = TestMatrixTensorHandle(ctx);
+  const char* device_type = TFE_TensorHandleDeviceType(hcpu, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_TRUE(absl::StrContains(device_type, "CPU")) << device_type;
+  int device_id = TFE_TensorHandleDeviceID(hcpu, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_EQ(0, device_id) << device_id;
+
+  // Disable the test if no GPU is present.
+  std::string gpu_device_name;
+  if (GetDeviceName(ctx, &gpu_device_name, "GPU")) {
+    TFE_TensorHandle* hgpu = TFE_TensorHandleCopyToDevice(
+        hcpu, ctx, gpu_device_name.c_str(), status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+    TFE_Op* shape_op = ShapeOp(ctx, hgpu);
+    TFE_OpSetDevice(shape_op, gpu_device_name.c_str(), status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+    TFE_TensorHandle* retvals[1];
+    int num_retvals = 1;
+    TFE_Execute(shape_op, &retvals[0], &num_retvals, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+    device_type = TFE_TensorHandleDeviceType(retvals[0], status.get());
+    ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+    ASSERT_TRUE(absl::StrContains(device_type, "GPU")) << device_type;
+
+    device_id = TFE_TensorHandleDeviceID(retvals[0], status.get());
+    ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+    ASSERT_EQ(0, device_id) << device_id;
+
+    TFE_DeleteOp(shape_op);
+    TFE_DeleteTensorHandle(retvals[0]);
+    TFE_DeleteTensorHandle(hgpu);
+  }
+
+  TFE_DeleteTensorHandle(hcpu);
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  TFE_DeleteExecutor(executor);
+  TFE_DeleteContext(ctx);
+}
+
+TEST(CAPI, TensorHandleDefaults) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status.get());
+  TFE_DeleteContextOptions(opts);
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+
+  TFE_TensorHandle* h_default = TestMatrixTensorHandle(ctx);
+  const char* device_type = TFE_TensorHandleDeviceType(h_default, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_TRUE(absl::StrContains(device_type, "CPU")) << device_type;
+  int device_id = TFE_TensorHandleDeviceID(h_default, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_EQ(0, device_id) << device_id;
+
+  TFE_TensorHandle* h_cpu = TFE_TensorHandleCopyToDevice(
+      h_default, ctx, "/device:CPU:0", status.get());
+  const char* device_type_cpu = TFE_TensorHandleDeviceType(h_cpu, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_TRUE(absl::StrContains(device_type_cpu, "CPU")) << device_type_cpu;
+  int device_id_cpu = TFE_TensorHandleDeviceID(h_cpu, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  ASSERT_EQ(0, device_id_cpu) << device_id_cpu;
+
+  TFE_DeleteTensorHandle(h_default);
+  TFE_DeleteTensorHandle(h_cpu);
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status.get());
+  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+  TFE_DeleteExecutor(executor);
+  TFE_DeleteContext(ctx);
+}
+
+TEST(CAPI, CreateLocalContextAsReset) {
+  tensorflow::ServerDef server_def = GetServerDef("worker", 2);
+  server_def.mutable_default_session_config()->set_isolate_session_state(false);
+
+  ServerFactory* factory;
+  ASSERT_TRUE(ServerFactory::GetFactory(server_def, &factory).ok());
+  server_def.set_job_name("worker");
+  server_def.set_task_index(0);
+  std::unique_ptr<tensorflow::ServerInterface> w0;
+  ASSERT_TRUE(
+      factory->NewServer(server_def, ServerFactory::Options(), &w0).ok());
+  ASSERT_TRUE(w0->Start().ok());
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::ServerInterface> w1;
+  ASSERT_TRUE(
+      factory->NewServer(server_def, ServerFactory::Options(), &w1).ok());
+  ASSERT_TRUE(w1->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  opts->session_options.options.config.set_isolate_session_state(false);
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  server_def.set_task_index(0);
+  auto cluster = server_def.mutable_cluster();
+  auto client_job = cluster->add_job();
+  client_job->set_name("localhost");
+  int client_port = tensorflow::testing::PickUnusedPortOrDie();
+  client_job->mutable_tasks()->insert(
+      {0, absl::StrCat("localhost:", client_port)});
+  server_def.set_job_name("localhost");
+  auto serialized = server_def.SerializeAsString();
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  server_def.set_job_name("worker");
+  server_def.set_task_index(0);
+  tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->mutable_job(0);
+  int worker_port = tensorflow::testing::PickUnusedPortOrDie();
+  job_def->mutable_tasks()->at(0) = absl::StrCat("localhost:", worker_port);
+  serialized = server_def.SerializeAsString();
+  TFE_InitializeLocalOnlyContext(ctx, 0, serialized.data(), serialized.size(),
+                                 status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_DeleteContextOptions(opts);
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+
+  w0.release();
+  w1.release();
+}
+
+TEST(CAPI, ShareVariableAcrossContextsAfterUpdateContextWorksWithTimeout) {
+  tensorflow::ServerDef server_def_0 = GetServerDef(3);
+  server_def_0.mutable_default_session_config()->set_isolate_session_state(
+      false);
+  tensorflow::ServerDef server_def_1 =
+      ReplaceTaskInServerDef(server_def_0, /*task_index=*/0);
+
+  // These server defs have task index set to 0.
+  std::string serialized_server_def_0 = server_def_0.SerializeAsString();
+  std::string serialized_server_def_1 = server_def_1.SerializeAsString();
+
+  // Create two worker tasks.
+  server_def_0.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def_0, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+  server_def_0.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def_0, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+
+  // Create two contexts.
+  int32_t init_timeout_in_ms = 300000;
+  TFE_Context* ctx_0 =
+      CreateContext(serialized_server_def_0,
+                    /*isolate_session_state=*/false, init_timeout_in_ms);
+  TFE_Context* ctx_1 =
+      CreateContext(serialized_server_def_1,
+                    /*isolate_session_state=*/false, init_timeout_in_ms);
+
+  // Remote device on `worker2`.
+  const char remote_device[] = "/job:localhost/replica:0/task:2/device:CPU:0";
+  // `ctx_0`, `ctx_1` contains `remote_device`.
+  {
+    const std::vector<std::string>& device_names = ListDeviceNames(ctx_0);
+    ASSERT_TRUE(std::find(device_names.begin(), device_names.end(),
+                          remote_device) != device_names.end());
+  }
+
+  {
+    const std::vector<std::string>& device_names = ListDeviceNames(ctx_1);
+    ASSERT_TRUE(std::find(device_names.begin(), device_names.end(),
+                          remote_device) != device_names.end());
+  }
+
+  // Create a variable using `ctx_0`.
+  // Replace worker1 using a new worker, and update the contexts.
+  // Read the variable using `ctx_1`. This read should succeed.
+  //
+  // 1. Create a variable on `remote_device`, using `ctx_0`.
+  TFE_TensorHandle* handle_0 =
+      CreateVariable(ctx_0, 1.2, remote_device, /*variable_name=*/"var");
+
+  // 2. Wait for `var` to be created and initialized on the worker.
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextAsyncWait(ctx_0, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+
+  int port = tensorflow::testing::PickUnusedPortOrDie();
+  // 3. Replace worker1 with a new worker in server_def_0 and server_def_1.
+  ReplaceTaskInServerDef(&server_def_0, /*task_index=*/1, "localhost", port);
+  ReplaceTaskInServerDef(&server_def_1, /*task_index=*/1, "localhost", port);
+  // 4. Start a new task to replace worker1.
+  server_def_0.set_task_index(1);
+  worker_server1.release();
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def_0, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+
+  // 5a. Update `ctx_0` with updated `server_def_0`.
+  {
+    server_def_0.set_task_index(0);
+    std::string serialized_update = server_def_0.SerializeAsString();
+    TF_Status* status = TF_NewStatus();
+    TFE_ContextUpdateServerDefWithTimeout(ctx_0, 0, serialized_update.data(),
+                                          serialized_update.size(),
+                                          init_timeout_in_ms, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TF_DeleteStatus(status);
+  }
+
+  // 5b. Update `ctx_1` with updated `server_def_1`.
+  {
+    server_def_1.set_task_index(0);
+    std::string serialized_update = server_def_1.SerializeAsString();
+    TF_Status* status = TF_NewStatus();
+    TFE_ContextUpdateServerDefWithTimeout(ctx_1, 0, serialized_update.data(),
+                                          serialized_update.size(),
+                                          init_timeout_in_ms, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TF_DeleteStatus(status);
+  }
+
+  // 6. Read `var` using `ctx_1`. This read should succeed since `ctx_1` was
+  // created with `isolate_session_state` set to false, and update should
+  // preserve it.
+  {
+    // Create a handle to `var`, using `ctx_1`.
+    TFE_TensorHandle* var_handle =
+        CreateVarHandle(ctx_1, remote_device, /*variable_name=*/"var");
+
+    TFE_TensorHandle* handle_1 = nullptr;
+    int num_retvals = 1;
+    TF_Status* status = TF_NewStatus();
+    TFE_Op* op = TFE_NewOp(ctx_1, "ReadVariableOp", status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+    TFE_OpAddInput(op, var_handle, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_Execute(op, &handle_1, &num_retvals, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    TFE_DeleteOp(op);
+
+    ASSERT_EQ(1, num_retvals);
+    EXPECT_EQ(TF_FLOAT, TFE_TensorHandleDataType(handle_1));
+    EXPECT_EQ(0, TFE_TensorHandleNumDims(handle_1, status));
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+    // Read the value of tensor handle `handle_1`.
+    float value = 0.0f;
+    TF_Tensor* t = TFE_TensorHandleResolve(handle_1, status);
+    ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+    ASSERT_EQ(sizeof(float), TF_TensorByteSize(t));
+    memcpy(&value, TF_TensorData(t), sizeof(float));
+    TF_DeleteTensor(t);
+    EXPECT_EQ(1.2f, value);
+    TFE_DeleteTensorHandle(handle_1);
+    TF_DeleteStatus(status);
+    TFE_DeleteTensorHandle(var_handle);
+  }
+
+  TFE_DeleteTensorHandle(handle_0);
+
+  TFE_DeleteContext(ctx_0);
+  TFE_DeleteContext(ctx_1);
+
+  worker_server1.release();
+  worker_server2.release();
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/eager/c_api_internal.h

@@ -0,0 +1,43 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_C_API_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_C_API_INTERNAL_H_
+
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/tfe_cancellation_manager_internal.h"  // IWYU pragma: export
+#include "tensorflow/c/eager/tfe_executor_internal.h"  // IWYU pragma: export
+#include "tensorflow/c/eager/tfe_monitoring_internal.h"  // IWYU pragma: export
+#include "tensorflow/c/eager/tfe_op_attrs_internal.h"  // IWYU pragma: export
+#include "tensorflow/c/eager/tfe_tensor_debug_info_internal.h"  // IWYU pragma: export
+
+// TODO(b/154564140): Move this to its own header. This requires splitting
+// c_api_experimental.h
+struct TFE_ContextOptions {
+  TF_SessionOptions session_options;
+  // true if async execution is enabled.
+  bool async = false;
+  TFE_ContextDevicePlacementPolicy device_placement_policy{
+      TFE_DEVICE_PLACEMENT_SILENT};
+  // If true, use TFRT backend
+  bool use_tfrt = false;
+  // Whether to run elementary eager ops wrapped in a call op.
+  bool run_eager_op_as_function = false;
+  // Whether to rewrite jit_compile functions.
+  bool jit_compile_rewrite = false;
+};
+
+#endif  // TENSORFLOW_C_EAGER_C_API_INTERNAL_H_

tensorflow/c/eager/c_api_remote_function_test.cc

@@ -0,0 +1,74 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_remote_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace {
+
+void TestRemoteExecuteSilentCopiesFunc(bool async, bool remote,
+                                       bool heavy_load_on_streaming_rpc,
+                                       bool remote_func_outputs = false,
+                                       bool has_packed_input = false) {
+  return TestRemoteExecuteSilentCopies(async, remote, /*func=*/true,
+                                       heavy_load_on_streaming_rpc,
+                                       remote_func_outputs, has_packed_input);
+}
+
+TEST(CAPI, RemoteExecuteSilentCopiesAsyncFunc) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/true,
+                                    /*heavy_load_on_streaming_rpc=*/false);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesFuncRemoteOutputs) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/false, /*remote=*/true,
+                                    /*heavy_load_on_streaming_rpc=*/false,
+                                    /*remote_func_outputs=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesAsyncFuncRemoteOutputs) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/true,
+                                    /*heavy_load_on_streaming_rpc=*/false,
+                                    /*remote_func_outputs=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocalAsyncFunc) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/false,
+                                    /*heavy_load_on_streaming_rpc=*/false);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocalFuncRemoteOutputs) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/false, /*remote=*/false,
+                                    /*heavy_load_on_streaming_rpc=*/false,
+                                    /*remote_func_outputs=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocalAsyncFuncRemoteOutputs) {
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/false,
+                                    /*heavy_load_on_streaming_rpc=*/false,
+                                    /*remote_func_outputs=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocalAsyncFuncOrdering) {
+  // A remote input may be not ready when we start running a function. Test that
+  // the function execution should wait until the remote input is ready.
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/false,
+                                    /*heavy_load_on_streaming_rpc=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesRemoteAsyncPackedInputFuncOrdering) {
+  // A remote input (packed) may be not ready when we start running a function.
+  // Test that the function execution should wait until the remote input is
+  // ready.
+  TestRemoteExecuteSilentCopiesFunc(/*async=*/true, /*remote=*/true,
+                                    /*heavy_load_on_streaming_rpc=*/true,
+                                    /*remote_func_outputs*/ true,
+                                    /*has_packed_input=*/true);
+}
+
+}  // namespace

tensorflow/c/eager/c_api_remote_test.cc

@@ -0,0 +1,140 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_remote_test_util.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/common_runtime/function_optimization_registry.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+namespace {
+
+using ::tensorflow::string;
+
+void TestRemoteExecute(bool async) {
+  tensorflow::ServerDef server_def = GetServerDef(2);
+
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+
+  std::unique_ptr<tensorflow::GrpcServer> worker_server;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server)
+                  .ok());
+  ASSERT_TRUE(worker_server->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts,
+                                             TFE_DEVICE_PLACEMENT_EXPLICIT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* h0_task0 = TestMatrixTensorHandle(ctx);
+  TFE_TensorHandle* h1_task0 = TestMatrixTensorHandle(ctx);
+  const char remote_device_name[] =
+      "/job:localhost/replica:0/task:1/device:CPU:0";
+  auto* h0_task1 =
+      TFE_TensorHandleCopyToDevice(h0_task0, ctx, remote_device_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  auto* h1_task1 =
+      TFE_TensorHandleCopyToDevice(h1_task0, ctx, remote_device_name, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
+  TFE_OpSetDevice(matmul, remote_device_name, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TFE_TensorHandle* retvals[1];
+  int num_retvals = 1;
+  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+
+  TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  float product[4] = {0};
+  EXPECT_EQ(sizeof(product), TF_TensorByteSize(t));
+  memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+  EXPECT_EQ(7, product[0]);
+  EXPECT_EQ(10, product[1]);
+  EXPECT_EQ(15, product[2]);
+  EXPECT_EQ(22, product[3]);
+
+  TFE_DeleteTensorHandle(h0_task0);
+  TFE_DeleteTensorHandle(h1_task0);
+  TFE_DeleteTensorHandle(h0_task1);
+  TFE_DeleteTensorHandle(h1_task1);
+  TFE_DeleteTensorHandle(retvals[0]);
+
+  TFE_DeleteOp(matmul);
+
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+  TFE_DeleteContext(ctx);
+
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server.release();
+}
+
+TEST(CAPI, RemoteExecute) { TestRemoteExecute(false); }
+TEST(CAPI, RemoteExecuteAsync) { TestRemoteExecute(true); }
+
+void TestRemoteExecuteSilentCopiesOp(bool async, bool remote,
+                                     bool remote_func_outputs = false) {
+  return TestRemoteExecuteSilentCopies(async, remote, /*func=*/false,
+                                       /*heavy_load_on_streaming_rpc=*/false,
+                                       remote_func_outputs);
+}
+
+TEST(CAPI, RemoteExecuteSilentCopies) {
+  TestRemoteExecuteSilentCopiesOp(/*async=*/false, /*remote=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesAsync) {
+  TestRemoteExecuteSilentCopiesOp(/*async=*/true, /*remote=*/true);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocal) {
+  TestRemoteExecuteSilentCopiesOp(/*async=*/false, /*remote=*/false);
+}
+TEST(CAPI, RemoteExecuteSilentCopiesLocalAsync) {
+  TestRemoteExecuteSilentCopiesOp(/*async=*/true, /*remote=*/false);
+}
+
+}  // namespace

tensorflow/c/eager/c_api_remote_test_util.cc

@@ -0,0 +1,236 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_remote_test_util.h"
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+using ::tensorflow::string;
+
+string MatMulFunction(const string& matmul_device) {
+  tensorflow::FunctionDef def;
+  CHECK(tensorflow::protobuf::TextFormat::ParseFromString(
+      absl::StrCat("    signature {"
+                   "      name: 'MatMulFunction'"
+                   "      input_arg {"
+                   "        name: 'a'"
+                   "        type: DT_FLOAT"
+                   "      }"
+                   "      input_arg {"
+                   "        name: 'b'"
+                   "        type: DT_FLOAT"
+                   "      }"
+                   "      output_arg {"
+                   "        name: 'm'"
+                   "        type: DT_FLOAT"
+                   "      }"
+                   "    }"
+                   "    node_def {"
+                   "      name: 'matmul'"
+                   "      op: 'MatMul'"
+                   "      input: 'a'"
+                   "      input: 'b'"
+                   "      device: '",
+                   matmul_device, "'",
+                   "      attr {"
+                   "        key: 'T'"
+                   "        value {"
+                   "          type: DT_FLOAT"
+                   "        }"
+                   "      }"
+                   "    }"
+                   "    ret {"
+                   "      key: 'm'"
+                   "      value: 'matmul:product'"
+                   "    }"),
+      &def));
+  return def.SerializeAsString();
+}
+
+void TestRemoteExecuteSilentCopies(bool async, bool remote, bool func,
+                                   bool heavy_load_on_streaming_rpc,
+                                   bool remote_func_outputs,
+                                   bool has_packed_input) {
+  CHECK(!has_packed_input || func);
+  tensorflow::ServerDef server_def = GetServerDef(3);
+
+  // This server def has the task index set to 0.
+  string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_TensorHandle* h0_task0 = TestMatrixTensorHandle(ctx);
+  TFE_TensorHandle* h1_task0 = TestMatrixTensorHandle(ctx);
+  std::vector<TFE_TensorHandle*> handles_task0;
+  if (heavy_load_on_streaming_rpc) {
+    // Send 50 tensor copy requests to simulate that there have been some RPC
+    // requests been enqueued.
+    for (int i = 0; i < 50; ++i) {
+      handles_task0.push_back(TestMatrixTensorHandle(ctx));
+    }
+  }
+  const char task1_name[] = "/job:localhost/replica:0/task:1/device:CPU:0";
+  const char task2_name[] = "/job:localhost/replica:0/task:2/device:CPU:0";
+
+  std::vector<TFE_TensorHandle*> handles_task2;
+  for (auto* h_task0 : handles_task0) {
+    handles_task2.push_back(
+        TFE_TensorHandleCopyToDevice(h_task0, ctx, task2_name, status));
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  }
+
+  auto* h1_task2 =
+      TFE_TensorHandleCopyToDevice(h1_task0, ctx, task2_name, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TFE_TensorHandle* packed_handle = nullptr;
+  if (has_packed_input) {
+    int num_replicas = 1;
+    std::vector<TFE_TensorHandle*> packed_handles = {h1_task2};
+    packed_handle = TFE_CreatePackedTensorHandle(ctx, packed_handles.data(),
+                                                 &num_replicas, status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  }
+
+  TFE_Op* matmul = nullptr;
+  if (func) {
+    const string matmul_device = remote_func_outputs ? task2_name : "";
+    string function_def = MatMulFunction(matmul_device);
+    TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(),
+                              status);
+    CHECK_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+    matmul = TFE_NewOp(ctx, "MatMulFunction", status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+    TFE_OpAddInput(matmul, h0_task0, status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+    TFE_OpAddInput(matmul, has_packed_input ? packed_handle : h1_task2, status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  } else {
+    // Handles are on task0 (local), and task2, but op is on task1.
+    matmul = MatMulOp(ctx, h0_task0, h1_task2);
+  }
+  if (remote) {
+    TFE_OpSetDevice(matmul, task1_name, status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  } else if (!async) {
+    // Set the local device to CPU to easily validate mirroring
+    string cpu_device_name;
+    ASSERT_TRUE(GetDeviceName(ctx, &cpu_device_name, "CPU"));
+    TFE_OpSetDevice(matmul, cpu_device_name.c_str(), status);
+    EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+    auto remote_arg =
+        tensorflow::TensorHandleFromInterface(tensorflow::unwrap(h1_task2));
+    // The input handles should never change since they have been mirrored.
+    ASSERT_FALSE(remote_arg->HasLocalMirror(nullptr));
+  }
+
+  TFE_TensorHandle* retvals[1];
+  int num_retvals = 1;
+  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  // TODO(gjn): Add support for waiting on async local mirrors
+  if (!remote && !async && !remote_func_outputs) {
+    auto remote_arg =
+        tensorflow::TensorHandleFromInterface(tensorflow::unwrap(h1_task2));
+    // The input handles should never change since they have been mirrored.
+    ASSERT_TRUE(remote_arg->HasLocalMirror(nullptr));
+  }
+
+  if (remote_func_outputs) {
+    const string backing_device =
+        TFE_TensorHandleBackingDeviceName(retvals[0], status);
+    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+    EXPECT_EQ(backing_device, task2_name);
+  }
+
+  auto* retval_task0 = TFE_TensorHandleCopyToDevice(
+      retvals[0], ctx, "/job:localhost/replica:0/task:0/device:CPU:0", status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+
+  TF_Tensor* t = TFE_TensorHandleResolve(retval_task0, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteTensorHandle(retval_task0);
+  float product[4] = {0};
+  EXPECT_EQ(sizeof(product), TF_TensorByteSize(t));
+  memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t));
+  TF_DeleteTensor(t);
+  EXPECT_EQ(7, product[0]);
+  EXPECT_EQ(10, product[1]);
+  EXPECT_EQ(15, product[2]);
+  EXPECT_EQ(22, product[3]);
+
+  TFE_DeleteTensorHandle(h0_task0);
+  TFE_DeleteTensorHandle(h1_task0);
+  if (packed_handle) {
+    TFE_DeleteTensorHandle(packed_handle);
+  }
+  TFE_DeleteTensorHandle(h1_task2);
+  TFE_DeleteTensorHandle(retvals[0]);
+  for (auto* h : handles_task0) {
+    TFE_DeleteTensorHandle(h);
+  }
+  for (auto* h : handles_task2) {
+    TFE_DeleteTensorHandle(h);
+  }
+
+  TFE_DeleteOp(matmul);
+
+  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
+  TFE_ExecutorWaitForAllPendingNodes(executor, status);
+  ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteExecutor(executor);
+  if (func) {
+    TFE_ContextRemoveFunction(ctx, "MatMulFunction", status);
+  }
+  TFE_DeleteContext(ctx);
+
+  TF_DeleteStatus(status);
+
+  // TODO(b/136478427): Figure out how to correctly shut the server down.
+  worker_server1.release();
+  worker_server2.release();
+}

tensorflow/c/eager/c_api_remote_test_util.h

@@ -0,0 +1,27 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_C_API_REMOTE_TEST_UTIL_H_
+#define TENSORFLOW_C_EAGER_C_API_REMOTE_TEST_UTIL_H_
+
+// Run a function containing a MatMul op and check its output.
+// If heavy_load_on_streaming_rpc is true, send some rpc requests before the one
+// which creates a remote input, to simulate a scenario that the remote input
+// is not ready when we start running an op or a function.
+void TestRemoteExecuteSilentCopies(bool async, bool remote, bool func,
+                                   bool heavy_load_on_streaming_rpc,
+                                   bool remote_func_outputs = false,
+                                   bool has_packed_input = false);
+
+#endif  // TENSORFLOW_C_EAGER_C_API_REMOTE_TEST_UTIL_H_

tensorflow/c/eager/c_api_test_util.cc

@@ -0,0 +1,603 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_test_util.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"
+#include "tensorflow/core/util/port.h"
+
+using tensorflow::string;
+using tensorflow::tstring;
+
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, float value) {
+  float data[] = {value};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, nullptr, 0, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx,
+                                         const tensorflow::tstring& value) {
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_STRING, nullptr, 0, status);
+  tstring* data = static_cast<tstring*>(TF_TensorData(t));
+  *data = value;
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, int value) {
+  int data[] = {value};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_INT32, nullptr, 0, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, bool value) {
+  bool data[] = {value};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_BOOL, nullptr, 0, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* DoubleTestMatrixTensorHandle(TFE_Context* ctx) {
+  int64_t dims[] = {2, 2};
+  double data[] = {1.0, 2.0, 3.0, 4.0};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_DOUBLE, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestMatrixTensorHandle(TFE_Context* ctx) {
+  int64_t dims[] = {2, 2};
+  float data[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestMatrixTensorHandleWithInput(TFE_Context* ctx,
+                                                  float data[], int64_t dims[],
+                                                  int num_dims) {
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t =
+      TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0], num_dims, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestTensorHandleWithDimsFloat(TFE_Context* ctx, float data[],
+                                                int64_t dims[], int num_dims) {
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t =
+      TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0], num_dims, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestTensorHandleWithDimsInt(TFE_Context* ctx, int data[],
+                                              int64_t dims[], int num_dims) {
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t =
+      TFE_AllocateHostTensor(ctx, TF_INT32, &dims[0], num_dims, status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestMatrixTensorHandle100x100(TFE_Context* ctx) {
+  constexpr int64_t dims[] = {100, 100};
+  constexpr int num_elements = dims[0] * dims[1];
+  float data[num_elements];
+  for (int i = 0; i < num_elements; ++i) {
+    data[i] = 1.0f;
+  }
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* DoubleTestMatrixTensorHandle3X2(TFE_Context* ctx) {
+  int64_t dims[] = {3, 2};
+  double data[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestMatrixTensorHandle3X2(TFE_Context* ctx) {
+  int64_t dims[] = {3, 2};
+  float data[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_TensorHandle* TestVariable(TFE_Context* ctx, float value,
+                               const std::string& device_name) {
+  TF_Status* status = TF_NewStatus();
+  // Create the variable handle.
+  TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
+  TFE_OpSetAttrString(op, "container", "localhost", 0);
+  TFE_OpSetAttrString(op, "shared_name", "", 0);
+  if (!device_name.empty()) {
+    TFE_OpSetDevice(op, device_name.c_str(), status);
+  }
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_TensorHandle* var_handle = nullptr;
+  int num_retvals = 1;
+  TFE_Execute(op, &var_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_DeleteOp(op);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  CHECK_EQ(1, num_retvals);
+
+  // Assign 'value' to it.
+  op = TFE_NewOp(ctx, "AssignVariableOp", status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  TFE_OpAddInput(op, var_handle, status);
+
+  // Convert 'value' to a TF_Tensor then a TFE_TensorHandle.
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> t(
+      TF_AllocateTensor(TF_FLOAT, nullptr, 0, sizeof(value)), TF_DeleteTensor);
+  memcpy(TF_TensorData(t.get()), &value, TF_TensorByteSize(t.get()));
+
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)>
+      value_handle(TFE_NewTensorHandle(t.get(), status),
+                   TFE_DeleteTensorHandle);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  TFE_OpAddInput(op, value_handle.get(), status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  num_retvals = 0;
+  TFE_Execute(op, nullptr, &num_retvals, status);
+  TFE_DeleteOp(op);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  CHECK_EQ(0, num_retvals);
+
+  TF_DeleteStatus(status);
+
+  return var_handle;
+}
+
+TFE_Op* AddOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "AddV2", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, a, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, b, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
+
+  return op;
+}
+
+TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "MatMul", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, a, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, b, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
+
+  return op;
+}
+
+TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "Identity", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, a, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
+
+  return op;
+}
+
+TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "Shape", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, a, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));
+
+  return op;
+}
+
+TFE_TensorHandle* TestAxisTensorHandle(TFE_Context* ctx) {
+  int64_t dims[] = {1};
+  int data[] = {1};
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_INT32, &dims[0],
+                                        sizeof(dims) / sizeof(int64_t), status);
+  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+TFE_Op* MinOp(TFE_Context* ctx, TFE_TensorHandle* input,
+              TFE_TensorHandle* axis) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "Min", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, input, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, axis, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpSetAttrBool(op, "keep_dims", 1);
+  TFE_OpSetAttrType(op, "Tidx", TF_INT32);
+  TF_DeleteStatus(status);
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(input));
+
+  return op;
+}
+
+TFE_Op* AllReduceOp(TFE_Context* ctx, TFE_TensorHandle* in, int group_size) {
+  TF_Status* status = TF_NewStatus();
+
+  TFE_Op* op = TFE_NewOp(ctx, "CollectiveReduce", status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, in, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(in));
+  TFE_OpSetAttrInt(op, "group_size", group_size);
+  TFE_OpSetAttrInt(op, "group_key", 123);
+  TFE_OpSetAttrInt(op, "instance_key", 456);
+  TFE_OpSetAttrString(op, "merge_op", "Add", 3);
+  TFE_OpSetAttrString(op, "final_op", "Id", 2);
+  std::vector<int64_t> subdiv_offsets;
+  TFE_OpSetAttrIntList(op, "subdiv_offsets", subdiv_offsets.data(),
+                       subdiv_offsets.size());
+
+  return op;
+}
+
+TFE_Op* SendOp(TFE_Context* ctx, TFE_TensorHandle* in,
+               const std::string& op_name, const std::string& send_device,
+               const std::string& recv_device,
+               uint64_t send_device_incarnation) {
+  TF_Status* status = TF_NewStatus();
+  TFE_Op* op = TFE_NewOp(ctx, op_name.c_str(), status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_OpAddInput(op, in, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+
+  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(in));
+  TFE_OpSetAttrString(op, "tensor_name", "dummy", 5);
+  TFE_OpSetAttrString(op, "send_device", send_device.c_str(),
+                      send_device.size());
+  TFE_OpSetAttrString(op, "recv_device", recv_device.c_str(),
+                      recv_device.size());
+  TFE_OpSetAttrInt(op, "send_device_incarnation", send_device_incarnation);
+
+  return op;
+}
+
+TFE_Op* RecvOp(TFE_Context* ctx, const std::string& op_name,
+               const std::string& send_device, const std::string& recv_device,
+               uint64_t send_device_incarnation) {
+  TF_Status* status = TF_NewStatus();
+  TFE_Op* op = TFE_NewOp(ctx, op_name.c_str(), status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteStatus(status);
+
+  TFE_OpSetAttrType(op, "tensor_type", TF_INT32);
+  TFE_OpSetAttrString(op, "tensor_name", "dummy", 5);
+  TFE_OpSetAttrString(op, "send_device", send_device.c_str(),
+                      send_device.size());
+  TFE_OpSetAttrString(op, "recv_device", recv_device.c_str(),
+                      recv_device.size());
+  TFE_OpSetAttrInt(op, "send_device_incarnation", send_device_incarnation);
+
+  return op;
+}
+
+bool GetDeviceName(TFE_Context* ctx, string* device_name,
+                   const char* device_type) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get());
+  CHECK_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+
+  const int num_devices = TF_DeviceListCount(devices);
+  for (int i = 0; i < num_devices; ++i) {
+    const string dev_type(TF_DeviceListType(devices, i, status.get()));
+    CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    const string dev_name(TF_DeviceListName(devices, i, status.get()));
+    CHECK_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    if (dev_type == device_type) {
+      *device_name = dev_name;
+      LOG(INFO) << "Found " << device_type << " device " << *device_name;
+      TF_DeleteDeviceList(devices);
+      return true;
+    }
+  }
+  TF_DeleteDeviceList(devices);
+  return false;
+}
+
+tensorflow::ServerDef GetServerDef(const string& job_name, int num_tasks) {
+  tensorflow::ServerDef server_def;
+  server_def.set_protocol("grpc");
+  server_def.set_job_name(job_name);
+  server_def.set_task_index(0);
+  tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->add_job();
+  job_def->set_name(job_name);
+  for (int i = 0; i < num_tasks; i++) {
+    int port = tensorflow::testing::PickUnusedPortOrDie();
+    job_def->mutable_tasks()->insert({i, absl::StrCat("localhost:", port)});
+    LOG(INFO) << "Picked test port: " << port << " for job: " << job_name
+              << ", task: " << i;
+  }
+  return server_def;
+}
+
+tensorflow::ServerDef GetServerDef(int num_tasks) {
+  return GetServerDef("localhost", num_tasks);
+}
+
+tensorflow::ServerDef GetMultiClientServerDef(const std::string& job_name,
+                                              int num_tasks,
+                                              int num_virtual_gpus) {
+  tensorflow::ServerDef server_def;
+  server_def.set_protocol("grpc");
+  server_def.set_job_name(job_name);
+  server_def.set_task_index(0);
+  tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->add_job();
+  job_def->set_name(job_name);
+  for (int i = 0; i < num_tasks; i++) {
+    int port = tensorflow::testing::PickUnusedPortOrDie();
+    job_def->mutable_tasks()->insert({i, absl::StrCat("localhost:", port)});
+  }
+  auto* config = server_def.mutable_default_session_config();
+  config->mutable_experimental()->set_collective_group_leader(
+      absl::StrCat("/job:", job_name, "/replica:0/task:", 0));
+  auto* rewrite_options =
+      config->mutable_graph_options()->mutable_rewrite_options();
+  rewrite_options->set_scoped_allocator_optimization(
+      tensorflow::RewriterConfig::ON);
+  rewrite_options->mutable_scoped_allocator_opts()->add_enable_op(
+      "CollectiveReduce");
+
+  if ((tensorflow::IsGoogleCudaEnabled() || tensorflow::IsBuiltWithROCm()) &&
+      num_virtual_gpus > 0) {
+    tensorflow::GPUOptions* gpu_options =
+        server_def.mutable_default_session_config()->mutable_gpu_options();
+    auto virtual_devices =
+        gpu_options->mutable_experimental()->add_virtual_devices();
+    for (int i = 0; i < num_virtual_gpus; ++i) {
+      virtual_devices->add_memory_limit_mb(200);
+    }
+  }
+  return server_def;
+}
+
+TFE_TensorHandle* CreateVarHandle(TFE_Context* ctx,
+                                  const std::string& device_name,
+                                  const std::string& variable_name) {
+  TF_Status* status = TF_NewStatus();
+  // Create the variable handle.
+  TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
+  TFE_OpSetAttrString(op, "container", "localhost", 0);
+  TFE_OpSetAttrString(op, "shared_name", variable_name.data(),
+                      variable_name.size());
+  if (!device_name.empty()) {
+    TFE_OpSetDevice(op, device_name.c_str(), status);
+  }
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_TensorHandle* var_handle = nullptr;
+  int num_retvals = 1;
+  TFE_Execute(op, &var_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_DeleteOp(op);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  CHECK_EQ(1, num_retvals);
+  TF_DeleteStatus(status);
+  return var_handle;
+}
+
+TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value,
+                                 const std::string& device_name,
+                                 const std::string& variable_name) {
+  TF_Status* status = TF_NewStatus();
+  TFE_TensorHandle* var_handle =
+      CreateVarHandle(ctx, device_name, variable_name);
+
+  // Assign 'value' to it.
+  TFE_Op* op = TFE_NewOp(ctx, "AssignVariableOp", status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  TFE_OpAddInput(op, var_handle, status);
+  if (!device_name.empty()) {
+    TFE_OpSetDevice(op, device_name.c_str(), status);
+  }
+
+  // Convert 'value' to a TF_Tensor then a TFE_TensorHandle.
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> t(
+      TF_AllocateTensor(TF_FLOAT, nullptr, 0, sizeof(value)), TF_DeleteTensor);
+  memcpy(TF_TensorData(t.get()), &value, TF_TensorByteSize(t.get()));
+
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)>
+      value_handle(TFE_NewTensorHandle(t.get(), status),
+                   TFE_DeleteTensorHandle);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  TFE_OpAddInput(op, value_handle.get(), status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  int num_retvals = 0;
+  TFE_Execute(op, nullptr, &num_retvals, status);
+  TFE_DeleteOp(op);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  CHECK_EQ(0, num_retvals);
+  TF_DeleteStatus(status);
+  return var_handle;
+}
+
+TFE_Context* CreateContext(const std::string& serialized_server_def,
+                           bool isolate_session_state,
+                           int64_t init_timeout_in_ms) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  opts->session_options.options.config.set_isolate_session_state(
+      isolate_session_state);
+  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(false));
+  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_ContextSetServerDefWithTimeout(ctx, 0, serialized_server_def.data(),
+                                     serialized_server_def.size(),
+                                     init_timeout_in_ms, status,
+                                     /*clear_existing_contexts=*/false);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+  TF_DeleteStatus(status);
+  return ctx;
+}
+
+tensorflow::ServerDef ReplaceTaskInServerDef(
+    const tensorflow::ServerDef& server_def, int task_index) {
+  tensorflow::ServerDef server_def_copy = server_def;
+  tensorflow::ClusterDef* cluster_def = server_def_copy.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->mutable_job(0);
+  const int port = tensorflow::testing::PickUnusedPortOrDie();
+  job_def->mutable_tasks()->at(task_index) = absl::StrCat("localhost:", port);
+  return server_def_copy;
+}
+
+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) = absl::StrCat(host, ":", port);
+}
+
+std::vector<std::string> ListDeviceNames(TFE_Context* ctx) {
+  TF_Status* status = TF_NewStatus();
+  std::vector<std::string> device_names;
+  TF_DeviceList* devices = TFE_ContextListDevices(ctx, status);
+  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  const int num_devices = TF_DeviceListCount(devices);
+  for (int i = 0; i < num_devices; ++i) {
+    device_names.emplace_back(TF_DeviceListName(devices, i, status));
+    EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
+  }
+  TF_DeleteDeviceList(devices);
+  TF_DeleteStatus(status);
+  return device_names;
+}

tensorflow/c/eager/c_api_test_util.h

@@ -0,0 +1,173 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_C_API_TEST_UTIL_H_
+#define TENSORFLOW_C_EAGER_C_API_TEST_UTIL_H_
+
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+// Return a tensor handle containing a float scalar
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, float value);
+
+// Return a tensor handle containing a int scalar
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, int value);
+
+// Return a tensor handle containing a bool scalar
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, bool value);
+
+// Return a tensor handle containing a tstring scalar
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx,
+                                         const tensorflow::tstring& value);
+
+// Return a tensor handle containing a 2x2 matrix of doubles
+TFE_TensorHandle* DoubleTestMatrixTensorHandle(TFE_Context* ctx);
+
+// Return a tensor handle containing a 2x2 matrix of floats
+TFE_TensorHandle* TestMatrixTensorHandle(TFE_Context* ctx);
+
+// Return a tensor handle containing 2D matrix containing given data and
+// dimensions
+TFE_TensorHandle* TestMatrixTensorHandleWithInput(TFE_Context* ctx,
+                                                  float data[], int64_t dims[],
+                                                  int num_dims);
+
+// Get a Matrix TensorHandle with given float values and dimensions
+TFE_TensorHandle* TestTensorHandleWithDimsFloat(TFE_Context* ctx, float data[],
+                                                int64_t dims[], int num_dims);
+
+// Get a Matrix TensorHandle with given int values and dimensions
+TFE_TensorHandle* TestTensorHandleWithDimsInt(TFE_Context* ctx, int data[],
+                                              int64_t dims[], int num_dims);
+
+// Return a tensor handle with given type, values and dimensions.
+template <class T, TF_DataType datatype>
+TFE_TensorHandle* TestTensorHandleWithDims(TFE_Context* ctx, const T* data,
+                                           const int64_t* dims, int num_dims) {
+  TF_Status* status = TF_NewStatus();
+  TF_Tensor* t = TFE_AllocateHostTensor(ctx, datatype, dims, num_dims, status);
+  memcpy(TF_TensorData(t), data, TF_TensorByteSize(t));
+  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TF_DeleteTensor(t);
+  TF_DeleteStatus(status);
+  return th;
+}
+
+// Return a scalar tensor handle with given values.
+template <class T, TF_DataType datatype>
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, const T value) {
+  T data[] = {value};
+  return TestTensorHandleWithDims<T, datatype>(ctx, data, nullptr, 0);
+}
+
+// Return a tensor handle containing a 100x100 matrix of floats
+TFE_TensorHandle* TestMatrixTensorHandle100x100(TFE_Context* ctx);
+
+// Return a tensor handle containing a 3x2 matrix of doubles
+TFE_TensorHandle* DoubleTestMatrixTensorHandle3X2(TFE_Context* ctx);
+
+// Return a tensor handle containing a 3x2 matrix of floats
+TFE_TensorHandle* TestMatrixTensorHandle3X2(TFE_Context* ctx);
+
+// Return a variable handle referring to a variable with the given initial value
+// on the given device.
+TFE_TensorHandle* TestVariable(TFE_Context* ctx, float value,
+                               const std::string& device_name = "");
+
+// Return an add op multiplying `a` by `b`.
+TFE_Op* AddOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);
+
+// Return a matmul op multiplying `a` by `b`.
+TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);
+
+// Return an identity op.
+TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);
+
+// Return a shape op fetching the shape of `a`.
+TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);
+
+// Return an allreduce op adding up input tensor `in` from `group_size` workers.
+TFE_Op* AllReduceOp(TFE_Context* ctx, TFE_TensorHandle* in, int group_size);
+
+// Return a SendOp op `op_name` with send input tensor `in` and attributes
+// `send_device`, `recv_device`, and `send_device_incarnation` set.
+TFE_Op* SendOp(TFE_Context* ctx, TFE_TensorHandle* in,
+               const std::string& op_name, const std::string& send_device,
+               const std::string& recv_device,
+               uint64_t send_device_incarnation);
+
+// Return a RecvOp op `op_name` with the attributes `send_device`,
+// `recv_device`, and `send_device_incarnation` set.
+TFE_Op* RecvOp(TFE_Context* ctx, const std::string& op_name,
+               const std::string& send_device, const std::string& recv_device,
+               uint64_t send_device_incarnation);
+
+// Return a 1-D INT32 tensor containing a single value 1.
+TFE_TensorHandle* TestAxisTensorHandle(TFE_Context* ctx);
+
+// Return an op taking minimum of `input` long `axis` dimension.
+TFE_Op* MinOp(TFE_Context* ctx, TFE_TensorHandle* input,
+              TFE_TensorHandle* axis);
+
+// If there is a device of type `device_type`, returns true
+// and sets 'device_name' accordingly.
+// `device_type` must be either "GPU" or "TPU".
+bool GetDeviceName(TFE_Context* ctx, std::string* device_name,
+                   const char* device_type);
+
+// Create a ServerDef with the given `job_name` and add `num_tasks` tasks in it.
+tensorflow::ServerDef GetServerDef(const std::string& job_name, int num_tasks);
+
+// Create a ServerDef with job name "localhost" and add `num_tasks` tasks in it.
+tensorflow::ServerDef GetServerDef(int num_tasks);
+
+// Create a multi-client ServerDef with the given `job_name`, add `num_tasks`
+// tasks and `num_virtual_gpus` virtual GPUs in it.
+tensorflow::ServerDef GetMultiClientServerDef(const std::string& job_name,
+                                              int num_tasks,
+                                              int num_virtual_gpus = 0);
+
+// Create a variable handle with name `variable_name` on a device with name
+// `device_name`.
+TFE_TensorHandle* CreateVarHandle(TFE_Context* ctx,
+                                  const std::string& device_name,
+                                  const std::string& variable_name);
+
+// Create a variable with value `value` and name `variable_name` on a device
+// with name `device_name`.
+TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value,
+                                 const std::string& device_name,
+                                 const std::string& variable_name);
+
+TFE_Context* CreateContext(const std::string& serialized_server_def,
+                           bool isolate_session_state,
+                           int64_t init_timeout_in_ms);
+
+tensorflow::ServerDef ReplaceTaskInServerDef(
+    const tensorflow::ServerDef& server_def, int task_index);
+
+void ReplaceTaskInServerDef(tensorflow::ServerDef* server_def, int task_index,
+                            const std::string& host, int port);
+
+std::vector<std::string> ListDeviceNames(TFE_Context* ctx);
+
+#endif  // TENSORFLOW_C_EAGER_C_API_TEST_UTIL_H_

tensorflow/c/eager/c_api_unified_experimental.cc

@@ -0,0 +1,250 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+using tensorflow::string;
+
+namespace tensorflow {
+namespace tracing {
+typedef absl::flat_hash_map<std::string, tracing::FactoryFunction> FactoriesMap;
+
+static FactoriesMap& GetFactories() {
+  static FactoriesMap* factories = new FactoriesMap;
+  return *factories;
+}
+
+static tracing::FactoryFunction default_factory;
+
+void RegisterTracingEngineFactory(const std::string& name,
+                                  FactoryFunction factory) {
+  assert((!GetFactories().count(name)) ||
+         (GetFactories()[name] == factory) &&
+             "Duplicate tracing factory registration");
+  GetFactories()[name] = factory;
+}
+
+absl::Status SetDefaultTracingEngine(const char* name) {
+  auto entry = GetFactories().find(name);
+  if (entry != GetFactories().end()) {
+    default_factory = GetFactories().find(name)->second;
+    return absl::OkStatus();
+  }
+  std::string msg = absl::StrCat(
+      "No tracing engine factory has been registered with the key '", name,
+      "' (available: ");
+  // Ensure deterministic (sorted) order in the error message
+  std::set<std::string> factories_sorted;
+  for (const auto& factory : GetFactories())
+    factories_sorted.insert(factory.first);
+  const char* comma = "";
+  for (const std::string& factory : factories_sorted) {
+    msg += comma + factory;
+    comma = ", ";
+  }
+  msg += ")";
+
+  return absl::InvalidArgumentError(msg);
+}
+
+static TracingContext* CreateTracingExecutionContext(const char* fn_name,
+                                                     TF_Status* s) {
+  if (default_factory) {
+    return default_factory(fn_name, s);
+  }
+  tsl::Set_TF_Status_from_Status(
+      s, absl::FailedPreconditionError("default_factory is nullptr"));
+  return nullptr;
+}
+
+}  // end namespace tracing
+}  // end namespace tensorflow
+
+// =============================================================================
+// Public C API entry points
+//
+// These are only the generic entry points for the C API. This file does not
+// have any visibility into the graph/eager implementation and is only providing
+// C bindings to the abstract classes defined in the
+// c_api_unified_experimental_internal.h header.
+//
+// =============================================================================
+
+using tensorflow::AbstractFunction;
+using tensorflow::AbstractTensorHandle;
+using tensorflow::DataType;
+using tensorflow::dyn_cast;
+using tensorflow::OutputList;
+using tensorflow::Status;
+using tensorflow::unwrap;
+using tensorflow::wrap;
+using tensorflow::tracing::CreateTracingExecutionContext;
+using tensorflow::tracing::SetDefaultTracingEngine;
+using tensorflow::tracing::TracingContext;
+using tensorflow::tracing::TracingOperation;
+using tensorflow::tracing::TracingTensorHandle;
+
+void TF_SetTracingImplementation(const char* name, TF_Status* s) {
+  tsl::Set_TF_Status_from_Status(s, SetDefaultTracingEngine(name));
+}
+
+// Creates a new TensorFlow function, it is an execution context attached to a
+// given tracing context.
+TF_ExecutionContext* TF_CreateFunction(const char* fn_name, TF_Status* s) {
+  return wrap(CreateTracingExecutionContext(fn_name, s));
+}
+
+TF_AbstractFunction* TF_FinalizeFunction(TF_ExecutionContext* ctx,
+                                         TF_OutputList* outputs, TF_Status* s) {
+  AbstractFunction* func;
+  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(ctx));
+  if (!tracing_ctx) {
+    tsl::Set_TF_Status_from_Status(
+        s, absl::InvalidArgumentError(
+               "Only TracingContext can be converted into a function."));
+    return nullptr;
+  }
+  tsl::Set_TF_Status_from_Status(s,
+                                 tracing_ctx->Finalize(unwrap(outputs), &func));
+  TF_DeleteExecutionContext(ctx);
+  return wrap(func);
+}
+
+TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
+                                           TF_DataType dtype, TF_Shape shape,
+                                           TF_Status* s) {
+  DCHECK_GE(shape.num_dims, -1);
+  TracingTensorHandle* t;
+  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(func));
+  if (!tracing_ctx) {
+    tsl::Set_TF_Status_from_Status(
+        s, absl::InvalidArgumentError(
+               "TF_AddFunctionParameter must be called on a TracingContext."));
+    return nullptr;
+  }
+  tensorflow::PartialTensorShape partial_shape;
+  if (shape.num_dims != -1) {
+    DCHECK(shape.dim_sizes != nullptr);
+    Status status = tensorflow::PartialTensorShape::MakePartialShape(
+        reinterpret_cast<int64_t*>(shape.dim_sizes), shape.num_dims,
+        &partial_shape);
+    if (!status.ok()) {
+      tsl::Set_TF_Status_from_Status(s, status);
+      return nullptr;
+    }
+  }
+  tsl::Set_TF_Status_from_Status(
+      s, tracing_ctx->AddParameter(static_cast<DataType>(dtype), partial_shape,
+                                   &t));
+  return wrap(t);
+}
+
+void TF_DeleteExecutionContext(TF_ExecutionContext* c) { unwrap(c)->Release(); }
+
+TF_AbstractOp* TF_NewAbstractOp(TF_ExecutionContext* c) {
+  return wrap((unwrap(c)->CreateOperation()));
+}
+
+void TF_DeleteAbstractOp(TF_AbstractOp* op) { unwrap(op)->Release(); }
+
+void TF_DeleteAbstractTensor(TF_AbstractTensor* t) { unwrap(t)->Unref(); }
+
+TF_OutputList* TF_NewOutputList() { return wrap(new OutputList); }
+void TF_DeleteOutputList(TF_OutputList* o) { delete unwrap(o); }
+void TF_OutputListSetNumOutputs(TF_OutputList* o, int num_outputs,
+                                TF_Status* s) {
+  unwrap(o)->expected_num_outputs = num_outputs;
+  unwrap(o)->outputs.clear();
+  unwrap(o)->outputs.resize(num_outputs);
+}
+int TF_OutputListNumOutputs(TF_OutputList* o) {
+  return unwrap(o)->outputs.size();
+}
+TF_AbstractTensor* TF_OutputListGet(TF_OutputList* o, int i) {
+  return wrap(unwrap(o)->outputs[i]);
+}
+void TF_OutputListPushBack(TF_OutputList* o, TF_AbstractTensor* tensor,
+                           TF_Status* s) {
+  unwrap(o)->outputs.push_back(unwrap(tensor));
+}
+
+void TF_AbstractOpSetOpType(TF_AbstractOp* op, const char* const op_type,
+                            TF_Status* s) {
+  tsl::Set_TF_Status_from_Status(
+      s, unwrap(op)->Reset(op_type,
+                           /*raw_device_name=*/nullptr));
+}
+
+void TF_AbstractOpSetOpName(TF_AbstractOp* op, const char* const op_name,
+                            TF_Status* s) {
+  TracingOperation* tracing_op = dyn_cast<TracingOperation>(unwrap(op));
+  if (!tracing_op) {
+    tsl::Set_TF_Status_from_Status(
+        s, absl::InvalidArgumentError(
+               "TF_AbstractOpSetOpName must be called on a TracingOperation."));
+    return;
+  }
+  tsl::Set_TF_Status_from_Status(s, tracing_op->SetOpName(op_name));
+}
+
+void TF_AbstractOpSetAttrType(TF_AbstractOp* op, const char* const attr_name,
+                              TF_DataType value, TF_Status* s) {
+  Status status =
+      unwrap(op)->SetAttrType(attr_name, static_cast<DataType>(value));
+  TF_SetStatus(s, static_cast<TF_Code>(status.code()),
+               absl::StatusMessageAsCStr(status));
+}
+
+void TF_ExecuteOperation(TF_AbstractOp* op, int num_inputs,
+                         TF_AbstractTensor* const* inputs, TF_OutputList* o,
+                         TF_Status* s) {
+  for (int i = 0; i < num_inputs; i++) {
+    tsl::Set_TF_Status_from_Status(s, unwrap(op)->AddInput(unwrap(inputs[i])));
+    if (TF_GetCode(s) != TF_OK) {
+      return;
+    }
+  }
+  int num_outputs = unwrap(o)->expected_num_outputs;
+  tsl::Set_TF_Status_from_Status(
+      s, unwrap(op)->Execute(
+             absl::MakeSpan(reinterpret_cast<AbstractTensorHandle**>(
+                                unwrap(o)->outputs.data()),
+                            unwrap(o)->outputs.size()),
+             &num_outputs));
+}
+
+void TF_DeleteAbstractFunction(TF_AbstractFunction* func) {
+  unwrap(func)->Unref();
+}
+
+void TF_ExecutionContextRegisterFunction(TF_ExecutionContext* ctx,
+                                         TF_AbstractFunction* func,
+                                         TF_Status* s) {
+  tsl::Set_TF_Status_from_Status(s,
+                                 unwrap(ctx)->RegisterFunction(unwrap(func)));
+}

tensorflow/c/eager/c_api_unified_experimental.h

@@ -0,0 +1,153 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_H_
+#define TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_H_
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// =============================================================================
+// Unified Execution APIs for Eager and tracing backends.
+// =============================================================================
+
+// -----------------------------------------------------------------------------
+// Core APIs
+// -----------------------------------------------------------------------------
+
+// A TF_ExecutionContext stores knowledge about how to execute an operation.
+// E.g. it could know whether we're in eager mode or graph mode, keeps track
+// of gradient tapes, etc.
+typedef struct TF_ExecutionContext TF_ExecutionContext;
+
+// A TF_AbstractTensor is an input to an operation. E.g. it could be a union
+// type of eager and graph tensors. It is also the result of executing an
+// operation.
+typedef struct TF_AbstractTensor TF_AbstractTensor;
+
+// A TF_AbstractOp is the metadata we need to execute an operation. E.g. this
+// could contain the op type and other attributes.
+typedef struct TF_AbstractOp TF_AbstractOp;
+
+// Stores a function representation that can be used for execution or for
+// setting functional attributes of other composite ops e.g. control flow.
+typedef struct TF_AbstractFunction TF_AbstractFunction;
+
+// This allows the client to swap the implementation of the tracing engine.
+// Any future call to TF_CreateFunction will use the implementation defined
+// here.
+void TF_SetTracingImplementation(const char* name, TF_Status*);
+
+// Creates a new TensorFlow function. A Function is an execution context, and as
+// such it can trace operations through TF_ExecuteOperation. After completing
+// tracing, a function can be obtained by TF_FinalizeFunction.
+TF_ExecutionContext* TF_CreateFunction(const char* fn_name, TF_Status* status);
+
+// Creates a context for eager execution of operations.
+TF_ExecutionContext* TF_NewEagerExecutionContext(TFE_ContextOptions*,
+                                                 TF_Status* s);
+void TF_DeleteExecutionContext(TF_ExecutionContext*);
+
+// Represents a (partially-defined) shape.
+typedef struct TF_Shape {
+  int num_dims;  // Must be >= -1; -1 represents unknown rank.
+  int64_t* dim_sizes;
+} TF_Shape;
+
+// Add a new parameter to a TensorFlow Function.
+TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
+                                           TF_DataType dtype, TF_Shape shape,
+                                           TF_Status* s);
+
+// Create an operation suitable to use with the provided context. The operation
+// requires its type (e.g. "AddV2") to be set independently.
+TF_AbstractOp* TF_NewAbstractOp(TF_ExecutionContext* ctx);
+void TF_DeleteAbstractOp(TF_AbstractOp*);
+
+// TODO(srbs): Add APIs for specifying attrs etc.
+// `op_type` must outlive `op`.
+void TF_AbstractOpSetOpType(TF_AbstractOp* op, const char* const op_type,
+                            TF_Status* s);
+// `op_name` must outlive `op`.
+void TF_AbstractOpSetOpName(TF_AbstractOp* op, const char* const op_name,
+                            TF_Status* s);
+// `attr_name` must outlive `op`.
+void TF_AbstractOpSetAttrType(TF_AbstractOp* op, const char* const attr_name,
+                              TF_DataType value, TF_Status* s);
+
+void TF_DeleteAbstractTensor(TF_AbstractTensor*);
+
+// TF_OutputList holds the list of TF_AbstractTensor that results from executing
+// an operation, or provided to create a function.
+// When executing an operation in an eager context, the expected number of
+// outputs must be set beforehand with `TF_OutputListSetNumOutputs`.
+typedef struct TF_OutputList TF_OutputList;
+TF_OutputList* TF_NewOutputList();
+void TF_DeleteOutputList(TF_OutputList* o);
+// Prepare tracing to the expected number of output for an operation.
+void TF_OutputListSetNumOutputs(TF_OutputList* o, int num_outputs, TF_Status*);
+// Return the number of outputs in the list.
+int TF_OutputListNumOutputs(TF_OutputList* o);
+// Return the `i`th output in the list.
+TF_AbstractTensor* TF_OutputListGet(TF_OutputList* o, int i);
+// Append a tensor at the end of the output list, growing its size by one.
+void TF_OutputListPushBack(TF_OutputList* o, TF_AbstractTensor* tensor,
+                           TF_Status*);
+
+// TF_ExecuteOperation will, if in eager mode, execute, if in graph mode, maybe
+// capture some inputs and then add a node in the graph. The output tensors are
+// returned through the provided TF_OutputList.
+// Any active tape will observe the effects of this execution.
+void TF_ExecuteOperation(TF_AbstractOp* op, int num_inputs,
+                         TF_AbstractTensor* const* inputs, TF_OutputList* o,
+                         TF_Status* s);
+
+// Creates a new TF_AbstractFunction from the current tracing states in the
+// context. The provided `ctx` is consumed by this API call and deleted.
+// The returned TF_AbstractFunction must be deleted by the client,
+// TODO(aminim): clarify the contract on the state of the context after this
+// call.
+TF_AbstractFunction* TF_FinalizeFunction(TF_ExecutionContext* ctx,
+                                         TF_OutputList*, TF_Status*);
+
+void TF_DeleteAbstractFunction(TF_AbstractFunction*);
+
+// Register the function with the given context. This is particularly useful for
+// making a function available to an eager context.
+void TF_ExecutionContextRegisterFunction(TF_ExecutionContext*,
+                                         TF_AbstractFunction*, TF_Status*);
+
+// -----------------------------------------------------------------------------
+// APIs specific to Eager modes
+// -----------------------------------------------------------------------------
+
+// Temporary APIs till we figure out how to create scalar valued Eager
+// tensors and how to get value out of eager abstract tensors.
+TF_AbstractTensor* TF_CreateAbstractTensorFromEagerTensor(TFE_TensorHandle* t,
+                                                          TF_Status* s);
+TFE_TensorHandle* TF_AbstractTensorGetEagerTensor(TF_AbstractTensor* at,
+                                                  TF_Status* s);
+TFE_Context* TF_ExecutionContextGetTFEContext(TF_ExecutionContext*,
+                                              TF_Status* s);
+
+#ifdef __cplusplus
+} /* end extern "C" */
+#endif
+
+#endif  // TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_H_

tensorflow/c/eager/c_api_unified_experimental_eager.cc

@@ -0,0 +1,82 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/strcat.h"
+
+// =============================================================================
+// Public C API entry points
+// These are only the entry points specific to the Eager API.
+// =============================================================================
+
+using tensorflow::AbstractContext;
+using tensorflow::AbstractTensorHandle;
+using tensorflow::dyn_cast;
+using tensorflow::ImmediateExecutionContext;
+using tensorflow::ImmediateExecutionTensorHandle;
+using tensorflow::string;
+using tensorflow::unwrap;
+using tensorflow::wrap;
+using tensorflow::strings::StrCat;
+
+TF_ExecutionContext* TF_NewEagerExecutionContext(TFE_ContextOptions* options,
+                                                 TF_Status* s) {
+  TFE_Context* c_ctx = TFE_NewContext(options, s);
+  if (TF_GetCode(s) != TF_OK) {
+    return nullptr;
+  }
+  return wrap(static_cast<AbstractContext*>(unwrap(c_ctx)));
+}
+
+TF_AbstractTensor* TF_CreateAbstractTensorFromEagerTensor(TFE_TensorHandle* t,
+                                                          TF_Status* s) {
+  return wrap(static_cast<AbstractTensorHandle*>(unwrap(t)));
+}
+
+TFE_TensorHandle* TF_AbstractTensorGetEagerTensor(TF_AbstractTensor* at,
+                                                  TF_Status* s) {
+  auto handle = dyn_cast<ImmediateExecutionTensorHandle>(unwrap(at));
+  if (!handle) {
+    string msg = absl::StrCat("Not an eager tensor handle.",
+                              reinterpret_cast<uintptr_t>(at));
+    TF_SetStatus(s, TF_INVALID_ARGUMENT, msg.c_str());
+    return nullptr;
+  }
+  return wrap(handle);
+}
+
+TFE_Context* TF_ExecutionContextGetTFEContext(TF_ExecutionContext* ctx,
+                                              TF_Status* s) {
+  auto imm_ctx = dyn_cast<ImmediateExecutionContext>(unwrap(ctx));
+  if (!imm_ctx) {
+    string msg =
+        absl::StrCat("Not an eager context.", reinterpret_cast<uintptr_t>(ctx));
+    TF_SetStatus(s, TF_INVALID_ARGUMENT, msg.c_str());
+    return nullptr;
+  }
+  return wrap(imm_ctx);
+}

tensorflow/c/eager/c_api_unified_experimental_graph.cc

@@ -0,0 +1,441 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/graph_function.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "xla/tsl/c/tsl_status_internal.h"
+#include "tensorflow/core/framework/full_type.pb.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/types.h"
+
+using tensorflow::dyn_cast;
+using tensorflow::string;
+using tensorflow::gtl::ArraySlice;
+
+namespace tensorflow {
+namespace tracing {
+namespace graph {
+
+class GraphContext;
+class GraphOperation;
+class GraphTensor;
+
+auto& kUnknownDim = shape_inference::InferenceContext::kUnknownDim;
+auto& kUnknownRank = shape_inference::InferenceContext::kUnknownRank;
+
+// GraphTensor wraps a `TF_Output`, i.e. a pointer to TF_Operation and the index
+// into the list of outputs for the operation.
+class GraphTensor : public TracingTensorHandle {
+ public:
+  explicit GraphTensor(TF_Output output, TF_Graph* graph)
+      : TracingTensorHandle(kGraph), output_(output), graph_(graph) {}
+
+  tensorflow::DataType DataType() const override {
+    return static_cast<tensorflow::DataType>(TF_OperationOutputType(output_));
+  }
+
+  absl::Status Shape(tensorflow::PartialTensorShape* shape) const override {
+    DCHECK(shape != nullptr);
+    TF_Status status;
+    int num_dims = TF_GraphGetTensorNumDims(graph_, output_, &status);
+    DCHECK_GE(num_dims, -1);
+    TF_RETURN_IF_ERROR(StatusFromTF_Status(&status));
+    if (num_dims == kUnknownRank) {
+      return absl::OkStatus();
+    }
+
+    std::vector<int64_t> dims(num_dims, kUnknownDim);
+    TF_GraphGetTensorShape(graph_, output_,
+                           reinterpret_cast<int64_t*>(dims.data()), num_dims,
+                           &status);
+    TF_RETURN_IF_ERROR(StatusFromTF_Status(&status));
+    TF_RETURN_IF_ERROR(tensorflow::TensorShapeUtils::MakeShape(dims, shape));
+
+    return absl::OkStatus();
+  }
+
+  tensorflow::FullTypeDef FullType() const override {
+    const FullTypeDef* ft;
+    mutex_lock l(graph_->mu);
+    graph_->graph.NodeType(output_.oper->node.name(), &ft);
+    if (ft == nullptr) {
+      return FullTypeDef();
+    } else {
+      return *ft;
+    }
+  }
+
+  TF_Output output_;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractTensorHandle* ptr) {
+    return ptr->getKind() == kGraph;
+  }
+
+ private:
+  TF_Graph* graph_;  // For shape inference.
+};
+
+// GraphOperation wraps and populates a TF_OperationDescription.
+class GraphOperation : public TracingOperation {
+ public:
+  explicit GraphOperation(TF_Graph* g) : TracingOperation(kGraph), g_(g) {}
+  void Release() override { delete this; }
+  absl::Status Reset(const char* op, const char* raw_device_name) override {
+    if (op_) {
+      return absl::FailedPreconditionError("Reset called on already built op.");
+    }
+    if (raw_device_name) {
+      device_name_ = raw_device_name;
+    }
+    op_type_ = op;
+    return absl::OkStatus();
+  }
+  absl::Status SetOpName(const char* const op_name) override {
+    if (op_) {
+      return absl::FailedPreconditionError(
+          "SetOpName called on already built op.");
+    }
+    if (op_type_.empty()) {
+      return absl::FailedPreconditionError(
+          "GraphOperation::Reset must be called before calling SetOpName.");
+    }
+    // TODO(b/145674566): We use Graph::NewName to get a unique name here but
+    // this may not be consistent with python's naming policy.
+    mutex_lock l(g_->mu);
+    op_.reset(new TF_OperationDescription(g_, op_type_.c_str(),
+                                          g_->graph.NewName(op_name).c_str()));
+    return absl::OkStatus();
+  }
+  const string& Name() const override { return op_type_; }
+  const string& DeviceName() const override { return device_name_; }
+
+  absl::Status SetDeviceName(const char* name) override {
+    // TODO(srbs): Implement this.
+    device_name_ = name;
+    return absl::OkStatus();
+  }
+
+  absl::Status AddInput(AbstractTensorHandle* input) override {
+    GraphTensor* t = dyn_cast<GraphTensor>(input);
+    if (!t) {
+      return absl::InvalidArgumentError("Unable to cast input to GraphTensor");
+    }
+    TF_AddInput(op_.get(), t->output_);
+    return absl::OkStatus();
+  }
+  absl::Status AddInputList(
+      absl::Span<AbstractTensorHandle* const> inputs) override {
+    std::vector<TF_Output> tf_outputs(inputs.size());
+    for (int i = 0; i < inputs.size(); i++) {
+      GraphTensor* t = dyn_cast<GraphTensor>(inputs[i]);
+      if (!t) {
+        return absl::InvalidArgumentError(
+            "Unable to cast input to GraphTensor");
+      }
+      tf_outputs[i] = t->output_;
+    }
+    TF_AddInputList(op_.get(), tf_outputs.data(), tf_outputs.size());
+    return absl::OkStatus();
+  }
+  absl::Status Execute(absl::Span<AbstractTensorHandle*> retvals,
+                       int* num_retvals) override {
+    auto* tf_opdesc = op_.release();
+    if (tf_opdesc == nullptr) {
+      return absl::InvalidArgumentError("AbstractOp is incomplete.");
+    }
+    TF_Status* s = TF_NewStatus();
+    auto* operation = TF_FinishOperation(tf_opdesc, s);
+    TF_RETURN_IF_ERROR(StatusFromTF_Status(s));
+    TF_DeleteStatus(s);
+    *num_retvals = TF_OperationNumOutputs(operation);
+    if (*num_retvals > retvals.size()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "retvals span capacity (", retvals.size(),
+          ") is smaller than operation output count (", *num_retvals, ")"));
+    }
+    for (int i = 0; i < *num_retvals; ++i) {
+      retvals[i] = new GraphTensor({operation, i}, g_);
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status SetAttrString(const char* attr_name, const char* data,
+                             size_t length) override {
+    absl::string_view s(data, length);
+    op_->node_builder.Attr(attr_name, s);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrInt(const char* attr_name, int64_t value) override {
+    op_->node_builder.Attr(attr_name, static_cast<int64_t>(value));
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrFloat(const char* attr_name, float value) override {
+    op_->node_builder.Attr(attr_name, value);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrBool(const char* attr_name, bool value) override {
+    op_->node_builder.Attr(attr_name, value);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrType(const char* const attr_name,
+                           DataType value) override {
+    if (!op_) {
+      return absl::Status(
+          absl::StatusCode::kFailedPrecondition,
+          "op_type and op_name must be specified before specifying attrs.");
+    }
+    op_->node_builder.Attr(attr_name, value);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrShape(const char* attr_name, const int64_t* dims,
+                            const int num_dims) override {
+    PartialTensorShape shape;
+    if (num_dims >= 0) {
+      shape = PartialTensorShape(absl::Span<const int64_t>(
+          reinterpret_cast<const int64_t*>(dims), num_dims));
+    }
+    op_->node_builder.Attr(attr_name, shape);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrFunction(const char* attr_name,
+                               const AbstractOperation* value) override {
+    return absl::UnimplementedError(
+        "SetAttrFunction has not been implemented yet.");
+  }
+  absl::Status SetAttrFunctionName(const char* attr_name, const char* value,
+                                   size_t length) override {
+    tensorflow::NameAttrList func_name;
+    func_name.set_name(string(value, value + length));
+    op_->node_builder.Attr(attr_name, func_name);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrTensor(const char* attr_name,
+                             AbstractTensorInterface* tensor) override {
+    return absl::UnimplementedError(
+        "SetAttrTensor has not been implemented yet.");
+  }
+  absl::Status SetAttrStringList(const char* attr_name,
+                                 const void* const* values,
+                                 const size_t* lengths,
+                                 int num_values) override {
+    if (strcmp(attr_name, tensorflow::kColocationAttrName) == 0) {
+      op_->colocation_constraints.clear();
+      for (int i = 0; i < num_values; ++i) {
+        op_->colocation_constraints.emplace(static_cast<const char*>(values[i]),
+                                            lengths[i]);
+      }
+    } else {
+      std::vector<absl::string_view> v;
+      v.reserve(num_values);
+      for (int i = 0; i < num_values; ++i) {
+        v.emplace_back(static_cast<const char*>(values[i]), lengths[i]);
+      }
+      op_->node_builder.Attr(attr_name, v);
+    }
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrFloatList(const char* attr_name, const float* values,
+                                int num_values) override {
+    op_->node_builder.Attr(attr_name,
+                           ArraySlice<const float>(values, num_values));
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrIntList(const char* attr_name, const int64_t* values,
+                              int num_values) override {
+    op_->node_builder.Attr(
+        attr_name, ArraySlice<const int64_t>(
+                       reinterpret_cast<const int64_t*>(values), num_values));
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrTypeList(const char* attr_name, const DataType* values,
+                               int num_values) override {
+    op_->node_builder.Attr(attr_name,
+                           ArraySlice<const DataType>(values, num_values));
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrBoolList(const char* attr_name,
+                               const unsigned char* values,
+                               int num_values) override {
+    std::unique_ptr<bool[]> b(new bool[num_values]);
+    for (int i = 0; i < num_values; ++i) {
+      b[i] = values[i];
+    }
+    op_->node_builder.Attr(attr_name,
+                           ArraySlice<const bool>(b.get(), num_values));
+
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrShapeList(const char* attr_name, const int64_t** dims,
+                                const int* num_dims, int num_values) override {
+    std::vector<PartialTensorShape> shapes;
+    shapes.reserve(num_values);
+    for (int i = 0; i < num_values; ++i) {
+      if (num_dims[i] < 0) {
+        shapes.emplace_back();
+      } else {
+        shapes.emplace_back(absl::Span<const int64_t>(
+            reinterpret_cast<const int64_t*>(dims[i]), num_dims[i]));
+      }
+    }
+    op_->node_builder.Attr(attr_name, shapes);
+    return absl::OkStatus();
+  }
+  absl::Status SetAttrFunctionList(
+      const char* attr_name,
+      absl::Span<const AbstractOperation*> values) override {
+    return absl::UnimplementedError(
+        "SetAttrFunctionList has not been implemented yet.");
+  }
+  // For LLVM style RTTI.
+  static bool classof(const AbstractOperation* ptr) {
+    return ptr->getKind() == kGraph;
+  }
+  ~GraphOperation() override {}
+
+ private:
+  friend class GraphContext;  // For access to op_.
+  TF_Graph* g_;
+  std::unique_ptr<TF_OperationDescription> op_;
+  // Hold `op_type` and `op_name` till both are available since we need both
+  // to build a graph operation.
+  string op_type_;
+  const char* op_name_ = nullptr;
+  // TODO(srbs): Use this.
+  string device_name_;
+};
+
+// GraphContext wraps a TF_Graph modeling a single function and manages the
+// "execution" of operation, i.e. adding them to the function.
+class GraphContext : public TracingContext {
+ public:
+  explicit GraphContext(const char* name)
+      : TracingContext(kGraph),
+        graph_(new TF_Graph(), TF_DeleteGraph),
+        name_(name) {}
+
+  void Release() override { delete this; }
+
+  TracingOperation* CreateOperation() override {
+    return new GraphOperation(graph_.get());
+  }
+
+  absl::Status AddParameter(DataType dtype, const PartialTensorShape& shape,
+                            TracingTensorHandle** output) override {
+    TracingOperationPtr operation(CreateOperation());
+    TF_RETURN_IF_ERROR(operation->Reset("Placeholder", nullptr));
+    TF_RETURN_IF_ERROR(
+        operation->SetOpName(absl::StrCat("_input_", inputs_.size()).c_str()));
+    TF_RETURN_IF_ERROR(operation->SetAttrType("dtype", dtype));
+    if (!shape.unknown_rank()) {
+      TF_RETURN_IF_ERROR(operation->SetAttrShape(
+          "shape", reinterpret_cast<int64_t*>(shape.dim_sizes().data()),
+          shape.dims()));
+    }
+    int num_outputs = 1;
+    std::vector<AbstractTensorHandle*> outputs(num_outputs);
+    TF_RETURN_IF_ERROR(operation->Execute(
+        absl::Span<AbstractTensorHandle*>(outputs), &num_outputs));
+
+    if (num_outputs != 1) {
+      return absl::InternalError(
+          absl::StrCat("Expected 1 output but found ", num_outputs));
+    }
+    auto* t = dyn_cast<GraphTensor>(outputs[0]);
+    if (!t) {
+      return absl::InvalidArgumentError("Unable to cast input to GraphTensor");
+    }
+    inputs_.push_back(t->output_);
+    *output = tensorflow::down_cast<TracingTensorHandle*>(outputs[0]);
+    return absl::OkStatus();
+  }
+
+  absl::Status Finalize(OutputList* outputs, AbstractFunction** f) override {
+    std::vector<TF_Output> graph_outputs;
+    graph_outputs.reserve(outputs->outputs.size());
+    for (auto* abstract_output : outputs->outputs) {
+      GraphTensor* output = dyn_cast<GraphTensor>(abstract_output);
+      if (!output) {
+        return absl::UnimplementedError(
+            "Returning a non-graph tensor from a function has not "
+            "been implemented yet.");
+      }
+      graph_outputs.push_back(output->output_);
+    }
+
+    auto s = TF_NewStatus();
+    auto func = TF_GraphToFunction(graph_.get(), name_.data(), 0, -1, nullptr,
+                                   inputs_.size(), inputs_.data(),
+                                   graph_outputs.size(), graph_outputs.data(),
+                                   nullptr, nullptr, name_.data(), s);
+    *f = new GraphFunction(std::move(func->record->fdef()));
+    TF_DeleteFunction(func);
+    TF_RETURN_IF_ERROR(StatusFromTF_Status(s));
+    TF_DeleteStatus(s);
+    return absl::OkStatus();
+  }
+
+  absl::Status RegisterFunction(AbstractFunction* func) override {
+    return absl::UnimplementedError(
+        "Registering graph functions has not been implemented yet.");
+  }
+
+  absl::Status RemoveFunction(const string& func) override {
+    return absl::UnimplementedError(
+        "GraphContext::RemoveFunction has not been implemented yet.");
+  }
+  // For LLVM style RTTI.
+  static bool classof(const AbstractContext* ptr) {
+    return ptr->getKind() == kGraph;
+  }
+
+ private:
+  std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)> graph_;
+  std::vector<TF_Output> inputs_;
+  string name_;
+};
+
+static TracingContext* GraphTracingFactory(const char* name, TF_Status* s) {
+  return new GraphContext(name);
+}
+
+// Register the tracing implemented in this file as the default tracing engine.
+static bool register_tracing = [] {
+  RegisterTracingEngineFactory("graphdef", GraphTracingFactory);
+  SetDefaultTracingEngine("graphdef").IgnoreError();
+  return true;
+}();
+
+}  // namespace graph
+}  // namespace tracing
+}  // namespace tensorflow

tensorflow/c/eager/c_api_unified_experimental_internal.h

@@ -0,0 +1,138 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_INTERNAL_H_
+
+#include <vector>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Represents the results of the execution of an operation.
+struct OutputList {
+  std::vector<AbstractTensorHandle*> outputs;
+  int expected_num_outputs = -1;
+};
+
+namespace tracing {
+
+// =============================================================================
+// Implementation detail for the unified execution APIs for Eager and tracing
+// backends (graph/MLIR).
+//
+// This defines a set of abstract classes that are intended to provide the
+// functionality of the opaque C types exposed in the public APIs defined in the
+// `c_api_unified_experimental.h` header.
+// =============================================================================
+
+// Represents either a MlirTensor or a GraphTensor.
+// This base class does not expose any public methods other than to distinguish
+// which subclass it actually is. The user is responsible to use the right
+// type of AbstractTensor in their context (do not pass an MlirTensor to a
+// GraphContext and vice-versa).
+class TracingTensorHandle : public AbstractTensorHandle {
+ protected:
+  explicit TracingTensorHandle(AbstractTensorHandleKind kind)
+      : AbstractTensorHandle(kind) {}
+
+ public:
+  // For LLVM style RTTI.
+  static bool classof(const AbstractTensorHandle* ptr) {
+    return ptr->getKind() == kGraph || ptr->getKind() == kMlir;
+  }
+};
+
+// An abstract operation describes an operation by its type, name, and
+// attributes. It can be "executed" by the context with some input tensors.
+// It is allowed to reusing the same abstract operation for multiple execution
+// on a given context, with the same or different input tensors.
+class TracingOperation : public AbstractOperation {
+ protected:
+  explicit TracingOperation(AbstractOperationKind kind)
+      : AbstractOperation(kind) {}
+
+ public:
+  // Sets the name of the operation: this is an optional identifier that is
+  // not intended to carry semantics and preserved/propagated without
+  // guarantees.
+  virtual absl::Status SetOpName(const char* op_name) = 0;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractOperation* ptr) {
+    return ptr->getKind() == kGraph || ptr->getKind() == kMlir;
+  }
+};
+
+namespace internal {
+struct TracingOperationDeleter {
+  void operator()(TracingOperation* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+}  // namespace internal
+
+using TracingOperationPtr =
+    std::unique_ptr<TracingOperation, internal::TracingOperationDeleter>;
+
+// This holds the context for the execution: dispatching operations either to an
+// MLIR implementation or to a graph implementation.
+class TracingContext : public AbstractContext {
+ protected:
+  explicit TracingContext(AbstractContextKind kind) : AbstractContext(kind) {}
+
+ public:
+  // Add a function parameter and return the corresponding tensor.
+  virtual absl::Status AddParameter(DataType dtype,
+                                    const PartialTensorShape& shape,
+                                    TracingTensorHandle**) = 0;
+
+  // Finalize this context and make a function out of it. The context is in a
+  // invalid state after this call and must be destroyed.
+  virtual absl::Status Finalize(OutputList* outputs, AbstractFunction**) = 0;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractContext* ptr) {
+    return ptr->getKind() == kGraph || ptr->getKind() == kMlir;
+  }
+};
+
+typedef TracingContext* (*FactoryFunction)(const char* fn_name, TF_Status*);
+absl::Status SetDefaultTracingEngine(const char* name);
+void RegisterTracingEngineFactory(const std::string& name,
+                                  FactoryFunction factory);
+}  // namespace tracing
+
+DEFINE_CONVERSION_FUNCTIONS(AbstractContext, TF_ExecutionContext)
+DEFINE_CONVERSION_FUNCTIONS(AbstractTensorHandle, TF_AbstractTensor)
+DEFINE_CONVERSION_FUNCTIONS(AbstractFunction, TF_AbstractFunction)
+DEFINE_CONVERSION_FUNCTIONS(AbstractOperation, TF_AbstractOp)
+DEFINE_CONVERSION_FUNCTIONS(OutputList, TF_OutputList)
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_C_API_UNIFIED_EXPERIMENTAL_INTERNAL_H_

tensorflow/c/eager/custom_device_test.cc

@@ -0,0 +1,387 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// A simple logging device to test custom device registration.
+#include <cstddef>
+#include <memory>
+
+#include "absl/strings/match.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/custom_device_testutil.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/platform/test.h"
+
+TEST(CUSTOM_DEVICE, RegisterSimpleDevice) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* context = TFE_NewContext(opts, status.get());
+  TFE_DeleteContextOptions(opts);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  bool arrived = false;
+  bool executed = false;
+  const char* name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  RegisterLoggingDevice(context, name, /*strict_scope_placement=*/true,
+                        &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_TensorHandle* hcpu = TestMatrixTensorHandle(context);
+  ASSERT_FALSE(arrived);
+  TFE_TensorHandle* hdevice =
+      TFE_TensorHandleCopyToDevice(hcpu, context, name, status.get());
+  ASSERT_TRUE(arrived);
+  ASSERT_FALSE(executed);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  size_t size = TFE_TensorHandleDeviceMemorySize(hdevice, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(4 * sizeof(float), size);
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> matmul(
+      MatMulOp(context, hcpu, hdevice), TFE_DeleteOp);
+  TFE_OpSetDevice(matmul.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_TensorHandle* retval;
+  int num_retvals = 1;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+
+  TFE_DeleteTensorHandle(retval);
+  TFE_DeleteTensorHandle(hcpu);
+  TFE_DeleteTensorHandle(hdevice);
+  TFE_DeleteContext(context);
+}
+
+TEST(CUSTOM_DEVICE, ResetOperation) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts, status.get()), TFE_DeleteContext);
+  TFE_DeleteContextOptions(opts);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  bool arrived = false;
+  bool executed = false;
+  const char* custom_device_name =
+      "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  RegisterLoggingDevice(context.get(), custom_device_name,
+                        /*strict_scope_placement=*/true, &arrived, &executed,
+                        status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> reused_op(
+      TFE_NewOp(context.get(), "Identity", status.get()), TFE_DeleteOp);
+  TFE_OpReset(reused_op.get(), "Identity", custom_device_name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(std::string(TFE_OpGetDevice(reused_op.get(), status.get())),
+            std::string(custom_device_name));
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpReset(reused_op.get(), "Identity",
+              "/job:localhost/replica:0/task:0/device:CPU:0", status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(std::string(TFE_OpGetDevice(reused_op.get(), status.get())),
+            std::string("/job:localhost/replica:0/task:0/device:CPU:0"));
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+}
+
+TEST(CUSTOM_DEVICE, MakeVariable) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  bool arrived = false;
+  bool executed = false;
+  const char* name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  RegisterLoggingDevice(context.get(), name, /*strict_scope_placement=*/true,
+                        &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Create a variable handle placed on the custom device.
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context.get(), "VarHandleOp", status.get()), TFE_DeleteOp);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(op.get(), "shape", {}, 0, status.get());
+  TFE_OpSetAttrString(op.get(), "container", "", 0);
+  TFE_OpSetAttrString(op.get(), "shared_name", "", 0);
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_TensorHandle* var_handle = nullptr;
+  int num_retvals = 1;
+  executed = false;
+  TFE_Execute(op.get(), &var_handle, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+  auto handle_cleaner = tensorflow::gtl::MakeCleanup(
+      [var_handle]() { TFE_DeleteTensorHandle(var_handle); });
+
+  // Assign to the variable, copying to the custom device.
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)> one(
+      TestScalarTensorHandle(context.get(), 111.f), TFE_DeleteTensorHandle);
+  op.reset(TFE_NewOp(context.get(), "AssignVariableOp", status.get()));
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  TFE_OpAddInput(op.get(), one.get(), status.get());
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  executed = false;
+  num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+
+  // Read the variable's value.
+  op.reset(TFE_NewOp(context.get(), "ReadVariableOp", status.get()));
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  TFE_OpSetDevice(op.get(), name, status.get());
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  executed = false;
+  num_retvals = 1;
+  TFE_TensorHandle* var_value = nullptr;
+  TFE_Execute(op.get(), &var_value, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+  auto value_cleaner = tensorflow::gtl::MakeCleanup(
+      [var_value]() { TFE_DeleteTensorHandle(var_value); });
+  ASSERT_EQ(std::string(name), std::string(TFE_TensorHandleBackingDeviceName(
+                                   var_value, status.get())));
+  TFE_TensorHandle* var_value_unpacked =
+      UnpackTensorHandle(var_value, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> resolved_value(
+      TFE_TensorHandleResolve(var_value_unpacked, status.get()),
+      TF_DeleteTensor);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(111., *static_cast<float*>(TF_TensorData(resolved_value.get())));
+
+  // Free the backing buffer for the variable.
+  op.reset(TFE_NewOp(context.get(), "DestroyResourceOp", status.get()));
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+}
+
+TEST(CUSTOM_DEVICE, AccessVariableOnCustomDevice) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  bool arrived = false;
+  bool executed = false;
+  const char* name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  RegisterLoggingDevice(context.get(), name, /*strict_scope_placement=*/false,
+                        &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Create a variable handle placed on the custom device.
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context.get(), "VarHandleOp", status.get()), TFE_DeleteOp);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(op.get(), "shape", {}, 0, status.get());
+  TFE_OpSetAttrString(op.get(), "container", "", 0);
+  TFE_OpSetAttrString(op.get(), "shared_name", "", 0);
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_TensorHandle* var_handle = nullptr;
+  int num_retvals = 1;
+  executed = false;
+  TFE_Execute(op.get(), &var_handle, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+  auto handle_cleaner = tensorflow::gtl::MakeCleanup(
+      [var_handle]() { TFE_DeleteTensorHandle(var_handle); });
+
+  // Assign to the variable, copying to the custom device.
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)> one(
+      TestScalarTensorHandle(context.get(), 111.f), TFE_DeleteTensorHandle);
+  op.reset(TFE_NewOp(context.get(), "AssignVariableOp", status.get()));
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  TFE_OpAddInput(op.get(), one.get(), status.get());
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  executed = false;
+  num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+
+  // Read the variable's value.
+  op.reset(TFE_NewOp(context.get(), "ReadVariableOp", status.get()));
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
+  executed = false;
+  num_retvals = 1;
+  TFE_TensorHandle* var_value = nullptr;
+  TFE_Execute(op.get(), &var_value, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+  ASSERT_EQ(std::string(name),
+            std::string(TFE_TensorHandleDeviceName(var_value, status.get())));
+  TFE_DeleteTensorHandle(var_value);
+
+  // Free the backing buffer for the variable.
+  op.reset(TFE_NewOp(context.get(), "DestroyResourceOp", status.get()));
+  TFE_OpAddInput(op.get(), var_handle, status.get());
+  TFE_OpSetDevice(op.get(), name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+}
+
+TEST(CUSTOM_DEVICE, InputBasedPlacement) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  const char* custom0 = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  const char* custom1 = "/job:localhost/replica:0/task:0/device:CUSTOM:1";
+  bool arrived = false;
+  bool executed = false;
+  RegisterLoggingDevice(context.get(), custom0,
+                        /*strict_scope_placement=*/false, &arrived, &executed,
+                        status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  RegisterLoggingDevice(context.get(), custom1,
+                        /*strict_scope_placement=*/true, &arrived, &executed,
+                        status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)> hcpu(
+      TestMatrixTensorHandle(context.get()), TFE_DeleteTensorHandle);
+  ASSERT_FALSE(arrived);
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)> hcustom0(
+      TFE_TensorHandleCopyToDevice(hcpu.get(), context.get(), custom0,
+                                   status.get()),
+      TFE_DeleteTensorHandle);
+  ASSERT_TRUE(arrived);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  arrived = false;
+  std::unique_ptr<TFE_TensorHandle, decltype(&TFE_DeleteTensorHandle)> hcustom1(
+      TFE_TensorHandleCopyToDevice(hcpu.get(), context.get(), custom1,
+                                   status.get()),
+      TFE_DeleteTensorHandle);
+  ASSERT_TRUE(arrived);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Base case: two CPU inputs executes fine.
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> matmul(
+      MatMulOp(context.get(), hcpu.get(), hcpu.get()), TFE_DeleteOp);
+  TFE_TensorHandle* retval;
+  int num_retvals = 1;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_DeleteTensorHandle(retval);
+
+  // Custom device: inputs in same custom device works.
+  matmul.reset(MatMulOp(context.get(), hcustom0.get(), hcustom0.get()));
+  num_retvals = 1;
+  executed = false;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  ASSERT_TRUE(executed);
+  TFE_DeleteTensorHandle(retval);
+
+  // Custom device: inputs in different custom devices fails.
+  matmul.reset(MatMulOp(context.get(), hcustom0.get(), hcustom1.get()));
+  num_retvals = 1;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  ASSERT_NE(TF_OK, TF_GetCode(status.get()));
+  ASSERT_TRUE(absl::StrContains(TF_Message(status.get()), custom0));
+  ASSERT_TRUE(absl::StrContains(TF_Message(status.get()), custom1));
+
+  // Custom device: mix of custom/physical places the op on the custom device.
+  matmul.reset(MatMulOp(context.get(), hcustom0.get(), hcpu.get()));
+  num_retvals = 1;
+  executed = false;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  EXPECT_TRUE(executed);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_DeleteTensorHandle(retval);
+
+  // Explicit placement still forces the op onto the requested device
+  matmul.reset(MatMulOp(context.get(), hcustom0.get(), hcpu.get()));
+  TFE_OpSetDevice(matmul.get(), "/job:localhost/replica:0/task:0/device:CPU:0",
+                  status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  num_retvals = 1;
+  executed = false;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  EXPECT_FALSE(executed);
+  ASSERT_FALSE(TF_GetCode(status.get()) == TF_OK);
+
+  // Custom devices can refuse to do type-based dispatch (as hcustom1 is
+  // configured to do)
+  matmul.reset(MatMulOp(context.get(), hcustom1.get(), hcpu.get()));
+  num_retvals = 1;
+  executed = false;
+  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
+  EXPECT_FALSE(executed);
+  ASSERT_FALSE(TF_GetCode(status.get()) == TF_OK);
+}
+
+TEST(CUSTOM_DEVICE, InvalidRegistrationError) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  bool arrived = false;
+  bool executed = false;
+  RegisterLoggingDevice(context.get(), "/device:CUSTOM:0",
+                        /*strict_scope_placement=*/true, &arrived, &executed,
+                        status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+      << TF_Message(status.get());
+
+  const char* name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  RegisterLoggingDevice(context.get(), name, /*strict_scope_placement=*/true,
+                        &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  RegisterLoggingDevice(context.get(), name, /*strict_scope_placement=*/true,
+                        &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_ALREADY_EXISTS)
+      << TF_Message(status.get());
+
+  RegisterLoggingDevice(
+      context.get(), "/job:localhost/replica:0/task:0/device:CPU:0",
+      /*strict_scope_placement=*/true, &arrived, &executed, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_ALREADY_EXISTS)
+      << TF_Message(status.get());
+}

tensorflow/c/eager/custom_device_testutil.cc

@@ -0,0 +1,213 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// A simple logging device to test custom device registration.
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
+
+namespace {
+
+struct LoggingDevice {
+  std::string device_name;
+  std::string underlying_device;
+  // Set to true whenever a TensorHandle is copied onto the device
+  bool* arrived_flag = nullptr;
+  // Set to true whenever an operation is executed
+  bool* executed_flag = nullptr;
+  // If true, only explicit op placements are accepted. If false, uses
+  // type-based dispatch.
+  bool strict_scope_placement = false;
+  // If true, the LoggingDevice owns the memory of arrived_flag and
+  // executed_flag, and is responsible for deleting them in DeleteLoggingDevice.
+  bool delete_flags = false;
+};
+
+struct LoggedTensor {
+  TFE_TensorHandle* tensor;
+  LoggedTensor() = delete;
+  explicit LoggedTensor(TFE_TensorHandle* tensor) : tensor(tensor) {}
+  ~LoggedTensor() { TFE_DeleteTensorHandle(tensor); }
+};
+
+int64_t LoggedTensorDim(void* data, int dim_index, TF_Status* status) {
+  return TFE_TensorHandleDim(reinterpret_cast<LoggedTensor*>(data)->tensor,
+                             dim_index, status);
+}
+
+int LoggedTensorNumDims(void* data, TF_Status* status) {
+  return TFE_TensorHandleNumDims(reinterpret_cast<LoggedTensor*>(data)->tensor,
+                                 status);
+}
+
+void LoggedTensorDeallocator(void* data) {
+  delete reinterpret_cast<LoggedTensor*>(data);
+}
+
+TFE_TensorHandle* MakeLoggedTensorHandle(TFE_Context* context,
+                                         const std::string& logging_device_name,
+                                         std::unique_ptr<LoggedTensor> t,
+                                         TF_Status* status) {
+  TF_DataType dtype = TFE_TensorHandleDataType(t->tensor);
+  TFE_CustomDeviceTensorHandleMethods handle_methods;
+  handle_methods.num_dims = &LoggedTensorNumDims;
+  handle_methods.dim = &LoggedTensorDim;
+  handle_methods.deallocator = &LoggedTensorDeallocator;
+  return TFE_NewCustomDeviceTensorHandle(context, logging_device_name.c_str(),
+                                         dtype, t.release(), handle_methods,
+                                         status);
+}
+
+TFE_TensorHandle* CopyToLoggingDevice(TFE_Context* context,
+                                      TFE_TensorHandle* tensor,
+                                      TF_Status* status, void* device_info) {
+  LoggingDevice* dev = reinterpret_cast<LoggingDevice*>(device_info);
+  TFE_TensorHandle* t = TFE_TensorHandleCopyToDevice(
+      tensor, context, dev->underlying_device.c_str(), status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  auto dst = std::make_unique<LoggedTensor>(t);
+  *(dev->arrived_flag) = true;
+  return MakeLoggedTensorHandle(context, dev->device_name, std::move(dst),
+                                status);
+}
+
+TFE_TensorHandle* CopyTensorFromLoggingDevice(TFE_Context* context,
+                                              TFE_TensorHandle* tensor,
+                                              const char* target_device_name,
+                                              TF_Status* status,
+                                              void* device_info) {
+  TF_SetStatus(status, TF_INTERNAL,
+               "Trying to copy a tensor out of a logging device.");
+  return nullptr;
+}
+
+void LoggingDeviceExecute(const TFE_Op* original_op, int* num_outputs,
+                          TFE_TensorHandle** outputs, TF_Status* s,
+                          void* device_info) {
+  const char* requested_placement = TFE_OpGetDevice(original_op, s);
+  if (TF_GetCode(s) != TF_OK) return;
+
+  LoggingDevice* dev = reinterpret_cast<LoggingDevice*>(device_info);
+  if (dev->strict_scope_placement && *requested_placement == '\0') {
+    TF_SetStatus(s, TF_INTERNAL,
+                 "Ops must be placed on the device explicitly, or their inputs "
+                 "first copied to other devices.");
+    return;
+  }
+  TFE_Context* context = TFE_OpGetContext(original_op, s);
+  if (TF_GetCode(s) != TF_OK) return;
+  const char* operation_name = TFE_OpGetName(original_op, s);
+  if (TF_GetCode(s) != TF_OK) return;
+  const TFE_OpAttrs* attributes = TFE_OpGetAttrs(original_op);
+
+  TFE_Op* raw_op = TFE_NewOp(context, operation_name, s);
+  if (TF_GetCode(s) != TF_OK) return;
+  auto op_cleanup =
+      tensorflow::gtl::MakeCleanup([raw_op]() { TFE_DeleteOp(raw_op); });
+  TFE_Op* op = raw_op;
+  TFE_OpAddAttrs(op, attributes);
+  TFE_OpSetDevice(op, dev->underlying_device.c_str(), s);
+  if (TF_GetCode(s) != TF_OK) return;
+  int num_inputs = TFE_OpGetFlatInputCount(original_op, s);
+  if (TF_GetCode(s) != TF_OK) return;
+  for (int j = 0; j < num_inputs; ++j) {
+    TFE_TensorHandle* input = TFE_OpGetFlatInput(original_op, j, s);
+    if (TF_GetCode(s) != TF_OK) return;
+    const char* input_device = TFE_TensorHandleDeviceName(input, s);
+    if (TF_GetCode(s) != TF_OK) return;
+    if (dev->device_name == input_device) {
+      LoggedTensor* t = reinterpret_cast<LoggedTensor*>(
+          TFE_TensorHandleDevicePointer(input, s));
+      if (TF_GetCode(s) != TF_OK) return;
+      TFE_OpAddInput(op, t->tensor, s);
+    } else {
+      TFE_OpAddInput(op, input, s);
+    }
+    if (TF_GetCode(s) != TF_OK) return;
+  }
+  std::vector<TFE_TensorHandle*> op_outputs(*num_outputs);
+  TFE_Execute(op, op_outputs.data(), num_outputs, s);
+  if (TF_GetCode(s) != TF_OK) return;
+  for (int i = 0; i < *num_outputs; ++i) {
+    auto logged_tensor = std::make_unique<LoggedTensor>(op_outputs[i]);
+    outputs[i] = MakeLoggedTensorHandle(context, dev->device_name,
+                                        std::move(logged_tensor), s);
+  }
+  *(dev->executed_flag) = true;
+}
+
+}  // namespace
+
+void DeleteLoggingDevice(void* device_info) {
+  LoggingDevice* dev = reinterpret_cast<LoggingDevice*>(device_info);
+  if (dev->delete_flags) {
+    delete dev->arrived_flag;
+    delete dev->executed_flag;
+  }
+  delete dev;
+}
+
+void RegisterLoggingDevice(TFE_Context* context, const char* name,
+                           bool strict_scope_placement, bool* arrived_flag,
+                           bool* executed_flag, TF_Status* status) {
+  TFE_CustomDevice custom_device;
+  custom_device.copy_tensor_to_device = &CopyToLoggingDevice;
+  custom_device.copy_tensor_from_device = &CopyTensorFromLoggingDevice;
+  custom_device.delete_device = &DeleteLoggingDevice;
+  custom_device.execute = &LoggingDeviceExecute;
+  LoggingDevice* device = new LoggingDevice;
+  device->arrived_flag = arrived_flag;
+  device->executed_flag = executed_flag;
+  device->device_name = name;
+  device->underlying_device = "/job:localhost/replica:0/task:0/device:CPU:0";
+  device->strict_scope_placement = strict_scope_placement;
+  TFE_RegisterCustomDevice(context, custom_device, name, device, status);
+}
+
+TFE_TensorHandle* UnpackTensorHandle(TFE_TensorHandle* logged_tensor_handle,
+                                     TF_Status* status) {
+  void* device_ptr =
+      TFE_TensorHandleDevicePointer(logged_tensor_handle, status);
+  if (TF_GetCode(status) != TF_OK || device_ptr == nullptr) {
+    return nullptr;
+  }
+  return reinterpret_cast<LoggedTensor*>(device_ptr)->tensor;
+}
+
+void AllocateLoggingDevice(const char* name, bool* arrived_flag,
+                           bool* executed_flag, TFE_CustomDevice** device,
+                           void** device_info) {
+  TFE_CustomDevice* custom_device = new TFE_CustomDevice;
+  custom_device->copy_tensor_to_device = &CopyToLoggingDevice;
+  custom_device->copy_tensor_from_device = &CopyTensorFromLoggingDevice;
+  custom_device->delete_device = &DeleteLoggingDevice;
+  custom_device->execute = &LoggingDeviceExecute;
+  *device = custom_device;
+  LoggingDevice* logging_device = new LoggingDevice;
+  logging_device->arrived_flag = arrived_flag;
+  logging_device->executed_flag = executed_flag;
+  logging_device->device_name = name;
+  logging_device->underlying_device =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  logging_device->strict_scope_placement = true;
+  logging_device->delete_flags = true;
+  *device_info = reinterpret_cast<void*>(logging_device);
+}

tensorflow/c/eager/custom_device_testutil.h

@@ -0,0 +1,37 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
+#define TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
+
+// A simple logging device to test custom device registration.
+#include <memory>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/tf_status.h"
+
+void RegisterLoggingDevice(TFE_Context* context, const char* name,
+                           bool strict_scope_placement, bool* arrived_flag,
+                           bool* executed_flag, TF_Status* status);
+void AllocateLoggingDevice(const char* name, bool* arrived_flag,
+                           bool* executed_flag, TFE_CustomDevice** device,
+                           void** device_info);
+TFE_TensorHandle* UnpackTensorHandle(TFE_TensorHandle* logged_tensor_handle,
+                                     TF_Status* status);
+void DeleteLoggingDevice(void* device_info);
+
+#endif  // TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_

tensorflow/c/eager/dlpack.cc

@@ -0,0 +1,409 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/dlpack.h"
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "include/dlpack/dlpack.h"  // from @dlpack
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_reference.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+namespace {
+
+// Managing context for the DLManagedTensor, will manage the lifetime of
+// DLManagedTensor. When calling DLManagedTensor::deleter, it will notify the
+// original framework of destruction, and this context will be deleted also.
+struct TfDlManagedTensorCtx {
+  TensorReference reference;
+  std::vector<int64_t> shape;
+  std::vector<int64_t> strides;
+  DLManagedTensor tensor;
+
+  explicit TfDlManagedTensorCtx(const TensorReference& ref) : reference(ref) {}
+};
+
+// Gets tensor from eager tensor handle.
+const Tensor* GetTensorFromHandle(TFE_TensorHandle* h, TF_Status* status) {
+  if (h == nullptr) {
+    status->status = absl::InvalidArgumentError("Invalid handle");
+    return nullptr;
+  }
+  tensorflow::ImmediateExecutionTensorHandle* unwrapped_handle =
+      tensorflow::unwrap(h);
+  if (!tensorflow::TensorHandle::classof(unwrapped_handle)) {
+    status->status = absl::InvalidArgumentError("Invalid handle");
+    return nullptr;
+  }
+  tensorflow::TensorHandle* handle =
+      tensorflow::TensorHandleFromInterface(unwrapped_handle);
+  if (handle->Type() != TensorHandle::LOCAL) {
+    status->status = absl::InvalidArgumentError(absl::StrCat(
+        "DLPack doesn't support ", handle->TypeString(), " tensor"));
+    return nullptr;
+  }
+  const tensorflow::Tensor* tensor;
+  status->status = handle->Tensor(&tensor);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+  return tensor;
+}
+
+// Deleter for DLManagedTensor
+void DLManagedTensorDeleter(DLManagedTensor* arg) {
+  TfDlManagedTensorCtx* owner =
+      static_cast<TfDlManagedTensorCtx*>(arg->manager_ctx);
+  owner->reference.Unref();
+  delete owner;
+}
+
+// Converts TF_DATAType to DLPack data type.
+DLDataType GetDlDataType(TF_DataType data_type, TF_Status* status) {
+  DLDataType dtype;
+  dtype.lanes = 1;
+  dtype.bits = TF_DataTypeSize(data_type) * 8;
+  switch (data_type) {
+    case TF_DataType::TF_BOOL:
+      dtype.code = DLDataTypeCode::kDLBool;
+      break;
+    case TF_DataType::TF_HALF:
+    case TF_DataType::TF_FLOAT:
+    case TF_DataType::TF_DOUBLE:
+      dtype.code = DLDataTypeCode::kDLFloat;
+      break;
+    case TF_DataType::TF_INT8:
+    case TF_DataType::TF_INT16:
+    case TF_DataType::TF_INT32:
+    case TF_DataType::TF_INT64:
+      dtype.code = DLDataTypeCode::kDLInt;
+      break;
+    case TF_DataType::TF_UINT8:
+    case TF_DataType::TF_UINT16:
+    case TF_DataType::TF_UINT32:
+    case TF_DataType::TF_UINT64:
+      dtype.code = DLDataTypeCode::kDLUInt;
+      break;
+    case TF_DataType::TF_BFLOAT16:
+      dtype.code = DLDataTypeCode::kDLBfloat;
+      break;
+    case TF_DataType::TF_COMPLEX64:
+    case TF_DataType::TF_COMPLEX128:
+      dtype.code = DLDataTypeCode::kDLComplex;
+      break;
+    default:
+      status->status = absl::InvalidArgumentError(
+          absl::StrCat(DataType_Name(static_cast<DataType>(data_type)),
+                       " is not supported by dlpack"));
+      break;
+  }
+  return dtype;
+}
+
+// Gets DLPack's DLDevice from eager tensor handle.
+DLDevice GetDlContext(TFE_TensorHandle* h, TF_Status* status) {
+  DLDevice ctx;
+  const char* device_name =
+      tensorflow::unwrap(h)->BackingDeviceName(&status->status);
+  DeviceNameUtils::ParsedName parsed_name;
+  tensorflow::DeviceNameUtils::ParseFullName(device_name, &parsed_name);
+  std::string device_type = parsed_name.type;
+  int device_id = 0;
+  if (parsed_name.has_id) {
+    device_id = parsed_name.id;
+  }
+
+  ctx.device_id = device_id;
+  if (device_type == "CPU") {
+    ctx.device_type = DLDeviceType::kDLCPU;
+  } else if (device_type == "GPU") {
+#if TENSORFLOW_USE_ROCM
+    ctx.device_type = DLDeviceType::kDLROCM;
+#else
+    ctx.device_type = DLDeviceType::kDLCUDA;
+#endif
+  } else {
+    status->status =
+        absl::InvalidArgumentError("Unsupported Device Type for dlpack");
+  }
+
+  return ctx;
+}
+
+// Converts DLDevice to TF device name.
+absl::optional<std::string> DeviceNameFromDlContext(const DLDevice& ctx,
+                                                    TF_Status* status) {
+  switch (ctx.device_type) {
+    case DLDeviceType::kDLCPU:
+      return "CPU:0";
+    case DLDeviceType::kDLCUDA:
+      return absl::StrCat("GPU:", ctx.device_id);
+    case DLDeviceType::kDLROCM:
+      return absl::StrCat("GPU:", ctx.device_id);
+    default:
+      return std::nullopt;
+  }
+}
+
+// Converts DLPack data type to TF_DATATYPE.
+absl::Status TfDataTypeFormDlDataType(const DLDataType& dtype,
+                                      TF_DataType* tf_dtype) {
+  switch (dtype.code) {
+    case DLDataTypeCode::kDLBool:
+      if (dtype.bits != 8) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("Only DLPack bools of bitwidth 8 are supported, got: ",
+                         dtype.bits));
+      }
+      *tf_dtype = TF_DataType::TF_BOOL;
+      return absl::OkStatus();
+
+    case DLDataTypeCode::kDLUInt:
+      switch (dtype.bits) {
+        case 8:
+          *tf_dtype = TF_DataType::TF_UINT8;
+          return absl::OkStatus();
+        case 16:
+          *tf_dtype = TF_DataType::TF_UINT16;
+          return absl::OkStatus();
+        case 32:
+          *tf_dtype = TF_DataType::TF_UINT32;
+          return absl::OkStatus();
+        case 64:
+          *tf_dtype = TF_DataType::TF_UINT64;
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(
+              absl::StrCat("Unsupported UInt bits: ", dtype.bits));
+      }
+      return absl::OkStatus();
+    case DLDataTypeCode::kDLInt:
+      switch (dtype.bits) {
+        case 8:
+          *tf_dtype = TF_DataType::TF_INT8;
+          return absl::OkStatus();
+        case 16:
+          *tf_dtype = TF_DataType::TF_INT16;
+          return absl::OkStatus();
+        case 32:
+          *tf_dtype = TF_DataType::TF_INT32;
+          return absl::OkStatus();
+        case 64:
+          *tf_dtype = TF_DataType::TF_INT64;
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(
+              absl::StrCat("Unsupported Int bits: ", dtype.bits));
+      }
+      return absl::OkStatus();
+    case DLDataTypeCode::kDLFloat:
+      switch (dtype.bits) {
+        case 16:
+          *tf_dtype = TF_DataType::TF_HALF;
+          return absl::OkStatus();
+        case 32:
+          *tf_dtype = TF_DataType::TF_FLOAT;
+          return absl::OkStatus();
+        case 64:
+          *tf_dtype = TF_DataType::TF_DOUBLE;
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(
+              absl::StrCat("Unsupported Float bits: ", dtype.bits));
+      }
+      break;
+    case DLDataTypeCode::kDLBfloat:
+      switch (dtype.bits) {
+        case 16:
+          *tf_dtype = TF_DataType::TF_BFLOAT16;
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(
+              absl::StrCat("Unsupported BFloat bits: ", dtype.bits));
+      }
+      break;
+    case DLDataTypeCode::kDLComplex:
+      switch (dtype.bits) {
+        case 64:
+          *tf_dtype = TF_DataType::TF_COMPLEX64;
+          return absl::OkStatus();
+        case 128:
+          *tf_dtype = TF_DataType::TF_COMPLEX128;
+          return absl::OkStatus();
+        default:
+          return absl::InvalidArgumentError(
+              absl::StrCat("Unsupported Complex bits: ", dtype.bits));
+      }
+      break;
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("Unsupported Type Codes: ", dtype.code));
+  }
+}
+
+// Wraps the deleter function of DLManagedTensor to match the function signature
+// TFE_NewTensorHandleFromDeviceMemory.
+void DeallocatorWrapperFunc(void* data, size_t len, void* dlmt_vptr) {
+  TFE_CallDLManagedTensorDeleter(dlmt_vptr);
+}
+
+// Checks whether the stride array matches the layout of compact, row-majored
+// data.
+bool IsValidStrideCompactRowMajorData(int64_t* shape_arr, int64_t* stride_arr,
+                                      int ndim) {
+  bool valid = true;
+  int64_t expected_stride = 1;
+  for (int i = ndim - 1; i >= 0; --i) {
+    // Empty tensors are always compact regardless of strides.
+    if (shape_arr[i] == 0) return true;
+    // Note that dimensions with size=1 can have any stride.
+    if (shape_arr[i] != 1 && stride_arr[i] != expected_stride) {
+      valid = false;
+    }
+    expected_stride *= shape_arr[i];
+  }
+  return valid;
+}
+}  // namespace
+
+void* TFE_GetDLDevice(TFE_TensorHandle* h, TF_Status* status) {
+  auto dl_device = GetDlContext(h, status);
+  return new DLDevice{dl_device.device_type, dl_device.device_id};
+}
+
+void TFE_CallDLManagedTensorDeleter(void* dlm_ptr) {
+  DLManagedTensor* dlMTensor = static_cast<DLManagedTensor*>(dlm_ptr);
+  if (dlMTensor->deleter != nullptr) {
+    dlMTensor->deleter(dlMTensor);
+  }
+}
+
+void* TFE_HandleToDLPack(TFE_TensorHandle* h, TF_Status* status) {
+  auto tf_dlm_context = GetDlContext(h, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  auto* tf_dlm_data = TFE_TensorHandleDevicePointer(h, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  const Tensor* tensor = GetTensorFromHandle(h, status);
+  if (!status->status.ok() || tensor == nullptr) {
+    return nullptr;
+  }
+  TF_DataType data_type = static_cast<TF_DataType>(tensor->dtype());
+
+  auto tf_dlm_type = GetDlDataType(data_type, status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
+  auto* tf_dlm_tensor_ctx = new TfDlManagedTensorCtx(tensor_ref);
+  tf_dlm_tensor_ctx->reference = tensor_ref;
+
+  DLManagedTensor* dlm_tensor = &tf_dlm_tensor_ctx->tensor;
+  dlm_tensor->manager_ctx = tf_dlm_tensor_ctx;
+  dlm_tensor->deleter = &DLManagedTensorDeleter;
+  dlm_tensor->dl_tensor.device = tf_dlm_context;
+  int ndim = tensor->dims();
+  dlm_tensor->dl_tensor.ndim = ndim;
+  dlm_tensor->dl_tensor.data = tf_dlm_data;
+  dlm_tensor->dl_tensor.dtype = tf_dlm_type;
+
+  std::vector<int64_t>* shape_arr = &tf_dlm_tensor_ctx->shape;
+  std::vector<int64_t>* stride_arr = &tf_dlm_tensor_ctx->strides;
+  shape_arr->resize(ndim);
+  stride_arr->resize(ndim, 1);
+  for (int i = 0; i < ndim; i++) {
+    (*shape_arr)[i] = tensor->dim_size(i);
+  }
+  for (int i = ndim - 2; i >= 0; --i) {
+    (*stride_arr)[i] = (*shape_arr)[i + 1] * (*stride_arr)[i + 1];
+  }
+
+  dlm_tensor->dl_tensor.shape = shape_arr->data();
+  // There are two ways to represent compact row-major data
+  // 1) nullptr indicates tensor is compact and row-majored.
+  // 2) fill in the strides array as the real case for compact row-major data.
+  // Here we choose option 2, since some frameworks didn't handle the strides
+  // argument properly.
+  dlm_tensor->dl_tensor.strides = stride_arr->data();
+
+  dlm_tensor->dl_tensor.byte_offset =
+      0;  // TF doesn't handle the strides and byte_offsets here
+  return static_cast<void*>(dlm_tensor);
+}
+
+TFE_TensorHandle* TFE_HandleFromDLPack(void* dlm, TF_Status* status,
+                                       TFE_Context* ctx) {
+  DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(dlm);
+  DLTensor* dl_tensor = &dlmt->dl_tensor;
+  absl::optional<std::string> device_name =
+      DeviceNameFromDlContext(dl_tensor->device, status);
+  if (!device_name.has_value()) {
+    status->status = absl::InvalidArgumentError("Unsupported Device Type");
+    return nullptr;
+  }
+  TF_DataType dtype;
+  absl::Status s = TfDataTypeFormDlDataType(dl_tensor->dtype, &dtype);
+  if (!s.ok()) {
+    status->status = std::move(s);
+    return nullptr;
+  }
+  int num_dims = dl_tensor->ndim;
+  const int64_t* dims = dl_tensor->shape;
+  void* data = dl_tensor->data;
+
+  if (dl_tensor->byte_offset != 0) {
+    status->status = absl::InvalidArgumentError(
+        absl::StrCat("Unsupported byte_offset (", dl_tensor->byte_offset,
+                     ") from DLPack, must be zero"));
+    return nullptr;
+  }
+
+  size_t total_bytes = dl_tensor->dtype.bits / 8;
+  for (int i = 0; i < num_dims; i++) {
+    total_bytes *= dims[i];
+  }
+
+  if (dl_tensor->strides != nullptr &&
+      !IsValidStrideCompactRowMajorData(dl_tensor->shape, dl_tensor->strides,
+                                        num_dims)) {
+    status->status =
+        absl::InvalidArgumentError("Invalid strides array from DLPack");
+    return nullptr;
+  }
+
+  TFE_TensorHandle* handle = TFE_NewTensorHandleFromDeviceMemory(
+      ctx, device_name.value().c_str(), dtype, dims, num_dims, data,
+      total_bytes, &DeallocatorWrapperFunc, dlmt, status);
+
+  return handle;
+}
+
+}  // namespace tensorflow

tensorflow/c/eager/dlpack.h

@@ -0,0 +1,47 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_DLPACK_H_
+#define TENSORFLOW_C_EAGER_DLPACK_H_
+
+#include "tensorflow/c/eager/c_api.h"
+
+namespace tensorflow {
+
+// PyCapsule name for DLPack Tensor
+const char* const kDlTensorCapsuleName = "dltensor";
+
+// Returns the DLDevice* for the given eager tensor handle.
+//
+// The caller takes ownership of the returned pointer and is responsible for
+// deleting it.
+TF_CAPI_EXPORT extern void* TFE_GetDLDevice(TFE_TensorHandle* h,
+                                            TF_Status* status);
+
+// Converts eager tensor handle to DLPack (DLManagedTensor*), and return the
+// void* for further PyCapsule construction.
+TF_CAPI_EXPORT extern void* TFE_HandleToDLPack(TFE_TensorHandle* h,
+                                               TF_Status* status);
+
+// Converts DLPack (DLManagedTensor*) to eager tensor handle.
+TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_HandleFromDLPack(void* dlm,
+                                                             TF_Status* status,
+                                                             TFE_Context* ctx);
+
+// Calls the destructor of DLManagedTensor, used in the destructor of PyCapsule.
+TF_CAPI_EXPORT extern void TFE_CallDLManagedTensorDeleter(void* dlm_ptr);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_DLPACK_H_

tensorflow/c/eager/dlpack_test.cc

@@ -0,0 +1,113 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/dlpack.h"
+
+#include <vector>
+
+#include "absl/strings/str_join.h"
+#include "include/dlpack/dlpack.h"  // from @dlpack
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+void TestHandleFromDLPack(TF_Status* status, TFE_Context* ctx,
+                          std::vector<int64_t> shape,
+                          std::vector<int64_t> strides) {
+  size_t num_elements = 1;
+  for (int i = 0; i < static_cast<int32_t>(shape.size()); ++i) {
+    num_elements *= shape[i];
+  }
+  std::vector<float> data(num_elements);
+  for (size_t j = 0; j < num_elements; ++j) {
+    data[j] = j;
+  }
+  DLManagedTensor dlm_in = {};
+  DLTensor* dltensor_in = &dlm_in.dl_tensor;
+  dltensor_in->data = data.data();
+  dltensor_in->device = {kDLCPU, 0};
+  dltensor_in->ndim = static_cast<int32_t>(shape.size());
+  dltensor_in->dtype = {kDLFloat, 32, 1};
+  dltensor_in->shape = shape.data();
+  dltensor_in->strides = strides.data();
+  TFE_TensorHandle* handle = TFE_HandleFromDLPack(&dlm_in, status, ctx);
+  ASSERT_NE(handle, nullptr)
+      << TF_Message(status) << " (shape=[" << absl::StrJoin(shape, ",")
+      << "], strides=[" << absl::StrJoin(strides, ",") << "])";
+
+  auto* dlm_out =
+      static_cast<DLManagedTensor*>(TFE_HandleToDLPack(handle, status));
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  const DLTensor* dltensor_out = &dlm_out->dl_tensor;
+  EXPECT_EQ(dltensor_out->device.device_type, dltensor_in->device.device_type);
+  EXPECT_EQ(dltensor_out->device.device_id, dltensor_in->device.device_id);
+  EXPECT_EQ(dltensor_out->ndim, dltensor_in->ndim);
+  EXPECT_EQ(dltensor_out->dtype.code, dltensor_in->dtype.code);
+  EXPECT_EQ(dltensor_out->dtype.bits, dltensor_in->dtype.bits);
+  EXPECT_EQ(dltensor_out->dtype.lanes, dltensor_in->dtype.lanes);
+  for (int i = 0; i < dltensor_in->ndim; ++i) {
+    EXPECT_EQ(dltensor_out->shape[i], dltensor_in->shape[i]);
+    if (dltensor_out->strides) {
+      if (i == dltensor_in->ndim - 1) {
+        EXPECT_EQ(dltensor_out->strides[i], 1);
+      } else {
+        EXPECT_EQ(dltensor_out->strides[i],
+                  dltensor_out->shape[i + 1] * dltensor_out->strides[i + 1]);
+      }
+    }
+  }
+  const float* data_in = static_cast<const float*>(dltensor_in->data);
+  const float* data_out = static_cast<const float*>(dltensor_out->data);
+  for (size_t j = 0; j < num_elements; ++j) {
+    EXPECT_EQ(data_out[j], data_in[j]);
+  }
+
+  TFE_CallDLManagedTensorDeleter(dlm_out);
+  TFE_DeleteTensorHandle(handle);
+}
+
+TEST(DLPack, HandleFromDLPackStrides) {
+  TF_Status* status = TF_NewStatus();
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_Context* ctx = TFE_NewContext(opts, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+  TFE_DeleteContextOptions(opts);
+
+  TestHandleFromDLPack(status, ctx, {}, {});
+  TestHandleFromDLPack(status, ctx, {4}, {});
+  TestHandleFromDLPack(status, ctx, {4}, {1});
+  TestHandleFromDLPack(status, ctx, {4, 3, 2}, {});
+  TestHandleFromDLPack(status, ctx, {4, 3, 2}, {6, 2, 1});
+  // Test that dims with size=1 can have any stride.
+  TestHandleFromDLPack(status, ctx, {1}, {1});
+  TestHandleFromDLPack(status, ctx, {1}, {0});
+  TestHandleFromDLPack(status, ctx, {4, 1, 2}, {2, 1, 1});
+  TestHandleFromDLPack(status, ctx, {4, 1, 2}, {2, 0, 1});
+  TestHandleFromDLPack(status, ctx, {4, 3, 1}, {3, 1, 1});
+  TestHandleFromDLPack(status, ctx, {4, 3, 1}, {3, 1, 0});
+  // Test that empty tensors can have any strides.
+  TestHandleFromDLPack(status, ctx, {4, 0, 2}, {0, 2, 1});
+  TestHandleFromDLPack(status, ctx, {4, 0, 2}, {0, 1, 1});
+  TestHandleFromDLPack(status, ctx, {4, 0, 2}, {0, 0, 1});
+  TestHandleFromDLPack(status, ctx, {4, 0, 2}, {0, 2, 0});
+
+  TFE_DeleteContext(ctx);
+  TF_DeleteStatus(status);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/eager/gradient_checker.cc

@@ -0,0 +1,220 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/gradient_checker.h"
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_tensor.h"
+
+namespace tensorflow {
+namespace gradients {
+
+using namespace std;
+
+// ================== Helper functions =================
+
+// Fills data with values [start,end) with given step size.
+void Range(vector<int32_t>* data, int32_t start, int32_t end,
+           int32_t step = 1) {
+  for (int32_t i = start; i < end; i += step) {
+    (*data)[i] = i;
+  }
+}
+
+// Fills out_dims with the dimensions of the given tensor.
+void GetDims(const TF_Tensor* t, int64_t* out_dims) {
+  int num_dims = TF_NumDims(t);
+  for (int i = 0; i < num_dims; i++) {
+    out_dims[i] = TF_Dim(t, i);
+  }
+}
+
+// Runs model as is if output is a scalar,
+// else sums the output tensor before returning.
+absl::Status RunAndMaybeSum(AbstractContext* ctx, Model forward,
+                            absl::Span<AbstractTensorHandle* const> inputs,
+                            absl::Span<AbstractTensorHandle*> outputs,
+                            bool use_function) {
+  AbstractTensorHandle* model_outputs[1];
+
+  // Run the model.
+  TF_RETURN_IF_ERROR(
+      RunModel(forward, ctx, inputs, model_outputs, use_function));
+  AbstractTensorHandlePtr model_out(model_outputs[0]);
+
+  TF_Tensor* model_out_tensor;
+  TF_RETURN_IF_ERROR(GetValue(model_out.get(), &model_out_tensor));
+  int num_dims_out = TF_NumDims(model_out_tensor);
+  TF_DeleteTensor(model_out_tensor);
+
+  // If the output is a scalar, then return the scalar output
+  if (num_dims_out == 0) {
+    outputs[0] = model_out.release();
+    return absl::OkStatus();
+  }
+
+  // Else, reduce sum the output to get a scalar
+
+  // Will sum all dimensions, so get a Tensor containing [0,...,num_dims_out-1].
+  AbstractTensorHandlePtr sum_dims;
+  {
+    vector<int32_t> vals(num_dims_out);
+    int64_t vals_shape[] = {num_dims_out};
+    Range(&vals, 0, num_dims_out);
+    AbstractTensorHandle* sum_dims_raw = nullptr;
+    TF_RETURN_IF_ERROR(TestTensorHandleWithDims<int32_t, TF_INT32>(
+        ctx, vals.data(), vals_shape, 1, &sum_dims_raw));
+    sum_dims.reset(sum_dims_raw);
+  }
+
+  // Reduce sum the output on all dimensions.
+  TF_RETURN_IF_ERROR(ops::Sum(ctx, model_out.get(), sum_dims.get(), &outputs[0],
+                              /*keep_dims=*/false, "sum_output"));
+  return absl::OkStatus();
+}
+// ========================= End Helper Functions==============================
+
+absl::Status CalcNumericalGrad(AbstractContext* ctx, Model forward,
+                               absl::Span<AbstractTensorHandle* const> inputs,
+                               int input_index, bool use_function,
+                               AbstractTensorHandle** numerical_grad) {
+  vector<AbstractTensorHandle*> theta_inputs(inputs.size());
+  for (int i{}; i < inputs.size(); ++i) {
+    theta_inputs[i] = inputs[i];
+  }
+
+  if (input_index < 0 || input_index >= inputs.size()) {
+    return absl::InvalidArgumentError("input_index out of bounds");
+  }
+  AbstractTensorHandle* theta =
+      theta_inputs[input_index];  // parameter we are grad checking
+
+  // Convert from AbstractTensor to TF_Tensor.
+  TF_Tensor* theta_tensor;
+  TF_RETURN_IF_ERROR(GetValue(theta, &theta_tensor));
+
+  if (TF_TensorType(theta_tensor) != TF_FLOAT) {
+    return absl::InvalidArgumentError("theta_tensor must be of type TF_FLOAT");
+  }
+
+  // Get number of elements and fill data.
+  int num_elems = TF_TensorElementCount(theta_tensor);
+  vector<float> theta_data(num_elems);
+  memcpy(theta_data.data(), TF_TensorData(theta_tensor),
+         TF_TensorByteSize(theta_tensor));
+
+  // Initialize space for the numerical gradient.
+  vector<float> dtheta_approx(num_elems);
+
+  // Get theta shape and store in theta_dims.
+  int num_dims = TF_NumDims(theta_tensor);
+  vector<int64_t> theta_dims(num_dims);
+  GetDims(theta_tensor, theta_dims.data());
+
+  // Initialize auxilary data structures.
+  vector<float> thetaPlus_data(num_elems);
+  vector<float> thetaMinus_data(num_elems);
+  AbstractTensorHandle* f_outputs[1];
+
+  // Numerical Grad Check
+  for (int i = 0; i < num_elems; i++) {
+    // Get relative epsilon value
+    float epsilon = theta_data[i] == 0 ? 1e-4 : std::abs(theta_data[i] * 1e-4);
+    AbstractTensorHandlePtr two_eps;
+    {
+      AbstractTensorHandle* two_eps_raw = nullptr;
+      TF_RETURN_IF_ERROR(TestScalarTensorHandle<float, TF_FLOAT>(
+          ctx, 2 * epsilon, &two_eps_raw));
+      two_eps.reset(two_eps_raw);
+    }
+
+    // Initialize theta[i] + epsilon.
+    memcpy(thetaPlus_data.data(), TF_TensorData(theta_tensor),
+           TF_TensorByteSize(theta_tensor));
+    thetaPlus_data[i] += epsilon;
+    AbstractTensorHandlePtr thetaPlus;
+    {
+      AbstractTensorHandle* thetaPlus_raw = nullptr;
+      TF_RETURN_IF_ERROR(TestTensorHandleWithDims<float, TF_FLOAT>(
+          ctx, thetaPlus_data.data(), theta_dims.data(), num_dims,
+          &thetaPlus_raw));
+      thetaPlus.reset(thetaPlus_raw);
+    }
+
+    // Initialize theta[i] - epsilon.
+    memcpy(&thetaMinus_data[0], TF_TensorData(theta_tensor),
+           TF_TensorByteSize(theta_tensor));
+    thetaMinus_data[i] -= epsilon;
+    AbstractTensorHandlePtr thetaMinus;
+    {
+      AbstractTensorHandle* thetaMinus_raw = nullptr;
+      TF_RETURN_IF_ERROR(TestTensorHandleWithDims<float, TF_FLOAT>(
+          ctx, thetaMinus_data.data(), theta_dims.data(), num_dims,
+          &thetaMinus_raw));
+      thetaMinus.reset(thetaMinus_raw);
+    }
+
+    // Get f(theta + eps):
+    theta_inputs[input_index] = thetaPlus.get();
+    TF_RETURN_IF_ERROR(
+        RunAndMaybeSum(ctx, forward, theta_inputs, f_outputs, use_function));
+    AbstractTensorHandlePtr fPlus(f_outputs[0]);
+
+    // Get f(theta - eps):
+    theta_inputs[input_index] = thetaMinus.get();
+    TF_RETURN_IF_ERROR(
+        RunAndMaybeSum(ctx, forward, theta_inputs, f_outputs, use_function));
+    AbstractTensorHandlePtr fMinus(f_outputs[0]);
+
+    // Take Difference of both estimates: (f(theta + eps) - f(theta - eps)).
+    TF_RETURN_IF_ERROR(
+        ops::Sub(ctx, fPlus.get(), fMinus.get(), f_outputs, "sub_top"));
+    AbstractTensorHandlePtr fDiff(f_outputs[0]);
+
+    // Calculate using the difference quotient definition:
+    // (f(theta + eps) - f(theta - eps)) / (2 * eps).
+    TF_RETURN_IF_ERROR(
+        ops::Div(ctx, fDiff.get(), two_eps.get(), f_outputs, "diff_quotient"));
+    AbstractTensorHandlePtr diff_quotient(f_outputs[0]);
+
+    TF_Tensor* grad_tensor;
+    TF_RETURN_IF_ERROR(GetValue(diff_quotient.get(), &grad_tensor));
+    if (TF_TensorType(grad_tensor) != TF_FLOAT) {
+      return absl::InvalidArgumentError("grad_tensor must be of type TF_FLOAT");
+    }
+    if (TF_TensorElementCount(grad_tensor) != 1) {
+      return absl::InvalidArgumentError("grad_tensor must be a scalar");
+    }
+    float grad_data[1];
+    memcpy(&grad_data[0], TF_TensorData(grad_tensor),
+           TF_TensorByteSize(grad_tensor));
+    TF_DeleteTensor(grad_tensor);
+    dtheta_approx[i] = grad_data[0];
+  }
+
+  // Populate *numerical_grad with the data from dtheta_approx.
+  TF_RETURN_IF_ERROR(TestTensorHandleWithDims<float, TF_FLOAT>(
+      ctx, dtheta_approx.data(), theta_dims.data(), num_dims, numerical_grad));
+  TF_DeleteTensor(theta_tensor);
+  return absl::OkStatus();
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/eager/gradient_checker.h

@@ -0,0 +1,46 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_GRADIENT_CHECKER_H_
+#define TENSORFLOW_C_EAGER_GRADIENT_CHECKER_H_
+
+#include <memory>
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+
+namespace tensorflow {
+namespace gradients {
+
+/* Returns numerical grad inside `dtheta_approx` given `forward` model and
+ * parameter specified by `input_index`.
+ *
+ * I.e. if y = <output of the forward model> and w = inputs[input_index],
+ * this will calculate dy/dw numerically.
+ *
+ * `use_function` indicates whether to use graph mode(true) or eager(false).
+ *
+ * `numerical_grad` is the pointer to the AbstractTensorHandle* which will
+ * hold the numerical gradient data at the end of the function.
+ */
+absl::Status CalcNumericalGrad(AbstractContext* ctx, Model forward,
+                               absl::Span<AbstractTensorHandle* const> inputs,
+                               int input_index, bool use_function,
+                               AbstractTensorHandle** numerical_grad);
+
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_GRADIENT_CHECKER_H_

tensorflow/c/eager/gradient_checker_test.cc

@@ -0,0 +1,238 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/gradient_checker.h"
+
+#include <cstdint>
+#include <cstring>
+#include <tuple>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+
+using tensorflow::TF_StatusPtr;
+
+void CompareNumericalAndManualGradients(
+    Model model, AbstractContext* ctx,
+    absl::Span<AbstractTensorHandle* const> inputs, int input_index,
+    float* expected_grad, int num_grad, bool use_function,
+    double abs_error = 1e-2) {
+  absl::Status s;
+  AbstractTensorHandlePtr numerical_grad;
+  {
+    AbstractTensorHandle* numerical_grad_raw;
+    s = CalcNumericalGrad(ctx, model, inputs, input_index, use_function,
+                          &numerical_grad_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    numerical_grad.reset(numerical_grad_raw);
+  }
+
+  TF_Tensor* numerical_tensor;
+  s = GetValue(numerical_grad.get(), &numerical_tensor);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  auto num_elem_numerical = TF_TensorElementCount(numerical_tensor);
+  ASSERT_EQ(num_elem_numerical, num_grad);
+
+  float* dnumerical = new float[num_elem_numerical]{0};
+  memcpy(&dnumerical[0], TF_TensorData(numerical_tensor),
+         TF_TensorByteSize(numerical_tensor));
+
+  for (int j = 0; j < num_grad; j++) {
+    ASSERT_NEAR(dnumerical[j], expected_grad[j], abs_error);
+  }
+  delete[] dnumerical;
+  TF_DeleteTensor(numerical_tensor);
+}
+
+absl::Status MatMulModel(AbstractContext* ctx,
+                         absl::Span<AbstractTensorHandle* const> inputs,
+                         absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0],
+                     /*transpose_a=*/false,
+                     /*transpose_b=*/false,
+                     /*grad_a=*/false, /*grad_b=*/false, "MatMul");
+}
+
+absl::Status MulModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Mul(ctx, inputs[0], inputs[1], &outputs[0], "Mul");
+}
+
+// TODO(vnvo2409): Add more tests from `python/ops/gradient_checker_v2_test.py`.
+// These tests should not be confused with `[*]_grad_test` which compare the
+// result of `gradient_checker` and `[*]_grad`. The tests here test the
+// functionality of `gradient_checker` by comparing the result with expected
+// manual user-provided gradients.
+class GradientCheckerTest
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+
+    {
+      absl::Status s = StatusFromTF_Status(status.get());
+      CHECK_EQ(errors::OK, s.code()) << s.message();
+    }
+
+    {
+      AbstractContext* ctx_raw = nullptr;
+      absl::Status s =
+          BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+      ASSERT_EQ(errors::OK, s.code()) << s.message();
+      ctx_.reset(ctx_raw);
+    }
+
+    // Computing numerical gradients with TensorFloat-32 is numerically
+    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
+    // low tolerances
+    enable_tensor_float_32_execution(false);
+  }
+
+  AbstractContextPtr ctx_;
+
+ public:
+  bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
+  bool UseFunction() const { return std::get<2>(GetParam()); }
+};
+
+TEST_P(GradientCheckerTest, TestMatMul) {
+  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  int64_t A_dims[] = {2, 2};
+  AbstractTensorHandlePtr A;
+  {
+    AbstractTensorHandle* A_raw;
+    absl::Status s = TestTensorHandleWithDims<float, TF_FLOAT>(
+        ctx_.get(), A_vals, A_dims, 2, &A_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    A.reset(A_raw);
+  }
+  float B_vals[] = {.5f, -1.0f, 1.0f, 1.0f};
+  int64_t B_dims[] = {2, 2};
+  AbstractTensorHandlePtr B;
+  {
+    AbstractTensorHandle* B_raw;
+    absl::Status s = TestTensorHandleWithDims<float, TF_FLOAT>(
+        ctx_.get(), B_vals, B_dims, 2, &B_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    B.reset(B_raw);
+  }
+
+  float expected_dA[4] = {-.5f, 2.0f, -.5f, 2.0f};
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndManualGradients(
+      MatMulModel, ctx_.get(), {A.get(), B.get()}, 0, expected_dA, 4,
+      UseFunction()));
+}
+
+TEST_P(GradientCheckerTest, TestMul) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandlePtr y;
+  {
+    AbstractTensorHandle* y_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx_.get(), 7.0f, &y_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    y.reset(y_raw);
+  }
+
+  float expected_dx[1] = {7.0f};
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndManualGradients(
+      MulModel, ctx_.get(), {x.get(), y.get()}, 0, expected_dx, 1,
+      UseFunction()));
+}
+
+TEST_P(GradientCheckerTest, InputIndexOutOfBounds) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx_.get(), 1.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandle* numerical_grad_raw;
+  // Test input_index = -1 with 1 input tensor.
+  absl::Status s =
+      CalcNumericalGrad(ctx_.get(), MulModel, {x.get()}, /*input_index=*/-1,
+                        UseFunction(), &numerical_grad_raw);
+  ASSERT_EQ(absl::StatusCode::kInvalidArgument, s.code());
+
+  // Test input_index = 1 with 1 input tensor.
+  s = CalcNumericalGrad(ctx_.get(), MulModel, {x.get()}, /*input_index=*/1,
+                        UseFunction(), &numerical_grad_raw);
+  ASSERT_EQ(absl::StatusCode::kInvalidArgument, s.code());
+}
+
+TEST_P(GradientCheckerTest, InvalidInputType) {
+  int32_t x_vals[] = {2};
+  int64_t x_dims[] = {1};
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw;
+    absl::Status s = TestTensorHandleWithDims<int32_t, TF_INT32>(
+        ctx_.get(), x_vals, x_dims, 1, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandle* numerical_grad_raw;
+  // Pass an INT32 tensor as input 0, expecting an error.
+  absl::Status s =
+      CalcNumericalGrad(ctx_.get(), MulModel, {x.get()}, /*input_index=*/0,
+                        UseFunction(), &numerical_grad_raw);
+  ASSERT_EQ(absl::StatusCode::kInvalidArgument, s.code());
+}
+
+#ifdef PLATFORM_GOOGLE
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, GradientCheckerTest,
+    ::testing::Combine(::testing::Values("graphdef"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#else
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, GradientCheckerTest,
+    ::testing::Combine(::testing::Values("graphdef"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#endif
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/eager/gradients.cc

@@ -0,0 +1,528 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/gradients.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/gradients_internal.h"
+#include "tensorflow/c/eager/tape.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/eager/attr_builder.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace {
+
+// TODO(b/172558015): Using the pointer address as the identifier for the tensor
+// may lead to collisions. Introduce another way to get a unique id for this
+// tensor.
+int64_t ToId(const AbstractTensorHandle* t) {
+  return static_cast<int64_t>(reinterpret_cast<uintptr_t>(t));
+}
+
+}  // namespace
+
+absl::Status GradientRegistry::Register(
+    const std::string& op_name,
+    GradientFunctionFactory gradient_function_factory) {
+  auto iter = registry_.find(op_name);
+  if (iter != registry_.end()) {
+    return absl::AlreadyExistsError(
+        absl::StrCat("Gradient already exists for op: ", op_name, "."));
+  }
+  registry_.insert({op_name, gradient_function_factory});
+  return absl::OkStatus();
+}
+absl::Status GradientRegistry::Lookup(
+    const ForwardOperation& op,
+    std::unique_ptr<GradientFunction>* gradient_function) const {
+  auto iter = registry_.find(op.op_name);
+  if (iter == registry_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("No gradient defined for op: ", op.op_name, "."));
+  }
+  gradient_function->reset(iter->second(op));
+  return absl::OkStatus();
+}
+
+TapeTensor::TapeTensor(AbstractTensorHandle* handle) : handle_(handle) {
+  if (handle_) {
+    handle_->Ref();
+  }
+}
+TapeTensor::TapeTensor(const TapeTensor& other) {
+  handle_ = other.handle_;
+  if (handle_) {
+    handle_->Ref();
+  }
+}
+TapeTensor& TapeTensor::operator=(const TapeTensor& other) {
+  if (this != &other) {
+    if (other.handle_) {
+      other.handle_->Ref();
+    }
+    if (handle_) {
+      handle_->Unref();
+    }
+    handle_ = other.handle_;
+  }
+  return *this;
+}
+TapeTensor::TapeTensor(TapeTensor&& other) noexcept : handle_(other.handle_) {
+  other.handle_ = nullptr;
+}
+TapeTensor& TapeTensor::operator=(TapeTensor&& other) noexcept {
+  if (this != &other) {
+    if (handle_) {
+      handle_->Unref();
+    }
+    handle_ = other.handle_;
+    other.handle_ = nullptr;
+  }
+  return *this;
+}
+TapeTensor::~TapeTensor() {
+  if (handle_) {
+    handle_->Unref();
+  }
+}
+
+int64_t TapeTensor::GetID() const { return ToId(handle_); }
+
+tensorflow::DataType TapeTensor::GetDType() const {
+  if (handle_ == nullptr) {
+    return DT_INVALID;
+  }
+  return handle_->DataType();
+}
+AbstractTensorHandle* TapeTensor::GetHandle() const { return handle_; }
+
+AbstractTensorHandle* TapeTensor::ZerosLike() const { return nullptr; }
+
+class TapeVSpace
+    : public eager::VSpace<AbstractTensorHandle, GradientFunction, TapeTensor> {
+ public:
+  explicit TapeVSpace(AbstractContext* ctx) : ctx_(ctx) {}
+  ~TapeVSpace() override = default;
+
+  // Returns the number of elements in the gradient tensor.
+  int64_t NumElements(AbstractTensorHandle* tensor) const override;
+
+  // Consumes references to the tensors in the gradient_tensors list and returns
+  // a tensor with the result.
+  AbstractTensorHandle* AggregateGradients(
+      absl::Span<AbstractTensorHandle* const> gradient_tensors) const override;
+
+  // Calls the passed-in backward function.
+  // op_type is the op's name provided in RecordOperation.
+  absl::Status CallBackwardFunction(
+      const std::string& op_type, GradientFunction* gradient_function,
+      const std::vector<int64_t>& unneeded_gradients,
+      absl::Span<AbstractTensorHandle* const> output_gradients,
+      absl::Span<AbstractTensorHandle*> result) const override;
+
+  // Builds a tensor filled with ones with the same shape and dtype as `t`.
+  absl::Status BuildOnesLike(const TapeTensor& t,
+                             AbstractTensorHandle** result) const override;
+
+  // Looks up the ID of a Gradient.
+  int64_t TensorId(AbstractTensorHandle* tensor) const override;
+
+  // Converts a Gradient to a TapeTensor.
+  TapeTensor TapeTensorFromGradient(AbstractTensorHandle* g) const override;
+
+  void MarkAsResult(AbstractTensorHandle* gradient) const override;
+
+  void DeleteGradient(AbstractTensorHandle* gradient) const override;
+
+ private:
+  // The context where the aggregation op `Add` is to be created.
+  AbstractContext* ctx_;
+};
+
+// Returns the number of elements in the gradient tensor.
+int64_t TapeVSpace::NumElements(AbstractTensorHandle* tensor) const {
+  // TODO(srbs): It seems like this is used only for performance optimization
+  // and not for correctness. The only downside of keeping this 1 seems to be
+  // that the gradient accumulation is unbounded and we will never
+  // aggressively aggregate accumulated gradients to recover memory.
+  // Revisit and fix.
+  return 1;
+}
+
+// Consumes references to the tensors in the gradient_tensors list and returns
+// a tensor with the result.
+AbstractTensorHandle* TapeVSpace::AggregateGradients(
+    absl::Span<AbstractTensorHandle* const> gradient_tensors) const {
+  if (gradient_tensors.size() == 1) {
+    return gradient_tensors[0];
+  }
+
+  AbstractOperationPtr op(ctx_->CreateOperation());
+  absl::Status s = op->Reset("AddN", /*raw_device_name=*/nullptr);
+  if (!s.ok()) {
+    return nullptr;
+  }
+  s = op->AddInputList(gradient_tensors);
+  if (!s.ok()) {
+    return nullptr;
+  }
+
+  int num_outputs = 1;
+  std::vector<AbstractTensorHandle*> outputs(num_outputs);
+  s = op->Execute(absl::Span<AbstractTensorHandle*>(outputs), &num_outputs);
+  if (!s.ok()) {
+    return nullptr;
+  }
+  return outputs[0];
+}
+
+// Calls the passed-in backward function.
+// op_type is the op's name provided in RecordOperation.
+absl::Status TapeVSpace::CallBackwardFunction(
+    const std::string& op_type, GradientFunction* gradient_function,
+    const std::vector<int64_t>& unneeded_gradients,
+    absl::Span<AbstractTensorHandle* const> output_gradients,
+    absl::Span<AbstractTensorHandle*> result) const {
+  if (gradient_function == nullptr) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Provided null gradient_function for '", op_type, "'.\n",
+        "If the intent is to treat this op as non-differentiable consider ",
+        "using RegisterNotDifferentiable or ",
+        "NotDifferentiableGradientFunction."));
+  }
+  return gradient_function->Compute(ctx_, output_gradients, result);
+}
+
+absl::Status TapeVSpace::BuildOnesLike(const TapeTensor& t,
+                                       AbstractTensorHandle** result) const {
+  AbstractOperationPtr op(ctx_->CreateOperation());
+  TF_RETURN_IF_ERROR(op->Reset("OnesLike", /*raw_device_name=*/nullptr));
+  if (isa<tracing::TracingOperation>(op.get())) {
+    TF_RETURN_IF_ERROR(dyn_cast<tracing::TracingOperation>(op.get())->SetOpName(
+        absl::StrCat("OnesLike", ToId(t.GetHandle())).c_str()));
+  }
+  TF_RETURN_IF_ERROR(op->AddInput(t.GetHandle()));
+  int num_outputs = 1;
+  std::vector<AbstractTensorHandle*> outputs(num_outputs);
+  TF_RETURN_IF_ERROR(
+      op->Execute(absl::Span<AbstractTensorHandle*>(outputs), &num_outputs));
+  *result = outputs[0];
+  return absl::OkStatus();
+}
+
+// Looks up the ID of a Gradient.
+int64_t TapeVSpace::TensorId(AbstractTensorHandle* tensor) const {
+  return ToId(tensor);
+}
+
+// Converts a Gradient to a TapeTensor.
+TapeTensor TapeVSpace::TapeTensorFromGradient(AbstractTensorHandle* g) const {
+  return TapeTensor(g);
+}
+
+void TapeVSpace::MarkAsResult(AbstractTensorHandle* gradient) const {}
+
+void TapeVSpace::DeleteGradient(AbstractTensorHandle* gradient) const {
+  gradient->Unref();
+}
+
+void Tape::Watch(const AbstractTensorHandle* t) {
+  GradientTape::Watch(ToId(t));
+}
+void Tape::RecordOperation(absl::Span<AbstractTensorHandle* const> inputs,
+                           absl::Span<AbstractTensorHandle* const> outputs,
+                           GradientFunction* gradient_function,
+                           const std::string& op_name) {
+  std::vector<int64_t> input_ids(inputs.size());
+  std::vector<tensorflow::DataType> input_dtypes(inputs.size());
+  for (int i = 0; i < inputs.size(); i++) {
+    input_ids[i] = ToId(inputs[i]);
+    input_dtypes[i] = inputs[i]->DataType();
+  }
+  std::vector<TapeTensor> tape_tensors;
+  tape_tensors.reserve(outputs.size());
+  for (auto t : outputs) {
+    tape_tensors.push_back(TapeTensor(t));
+  }
+  GradientTape::RecordOperation(
+      op_name, tape_tensors, input_ids, input_dtypes,
+      [gradient_function]() -> GradientFunction* { return gradient_function; },
+      [](GradientFunction* ptr) {
+        if (ptr) {
+          delete ptr;
+        }
+      });
+}
+bool Tape::ShouldRecord(
+    absl::Span<const AbstractTensorHandle* const> tensors) const {
+  std::vector<int64_t> tensor_ids(tensors.size());
+  std::vector<tensorflow::DataType> tensor_dtypes(tensors.size());
+  for (int i = 0; i < tensors.size(); i++) {
+    tensor_ids[i] = ToId(tensors[i]);
+    tensor_dtypes[i] = tensors[i]->DataType();
+  }
+  return GradientTape::ShouldRecord(tensor_ids, tensor_dtypes);
+}
+void Tape::DeleteTrace(const AbstractTensorHandle* t) {
+  GradientTape::DeleteTrace(ToId(t));
+}
+
+std::vector<int64_t> MakeTensorIDList(
+    absl::Span<AbstractTensorHandle* const> tensors) {
+  std::vector<int64_t> ids(tensors.size());
+  for (int i = 0; i < tensors.size(); i++) {
+    ids[i] = ToId(tensors[i]);
+  }
+  return ids;
+}
+
+absl::Status Tape::ComputeGradient(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> targets,
+    absl::Span<AbstractTensorHandle* const> sources,
+    absl::Span<AbstractTensorHandle* const> output_gradients,
+    absl::Span<AbstractTensorHandle*> result) {
+  TapeVSpace vspace(ctx);
+  std::vector<int64_t> target_tensor_ids = MakeTensorIDList(targets);
+  std::vector<int64_t> source_tensor_ids = MakeTensorIDList(sources);
+  tensorflow::gtl::FlatSet<int64_t> sources_set(source_tensor_ids.begin(),
+                                                source_tensor_ids.end());
+  std::unordered_map<int64_t, TapeTensor> sources_that_are_targets;
+  for (int i = 0; i < target_tensor_ids.size(); ++i) {
+    int64_t target_id = target_tensor_ids[i];
+    if (sources_set.find(target_id) != sources_set.end()) {
+      auto tensor = targets[i];
+      sources_that_are_targets.insert(
+          std::make_pair(target_id, TapeTensor(tensor)));
+    }
+  }
+
+  TF_RETURN_IF_ERROR(GradientTape::ComputeGradient(
+      vspace, target_tensor_ids, source_tensor_ids, sources_that_are_targets,
+      output_gradients, result, /*build_default_zeros_grads*/ false));
+  return absl::OkStatus();
+}
+
+// Helper functions which delegate to `AbstractOperation`, update
+// the state of the ForwardOperation and call the tape as appropriate.
+// These APIs are mainly to facilitate testing and are subject to change.
+namespace internal {
+absl::Status Reset(AbstractOperation* op_, const char* op,
+                   const char* raw_device_name, ForwardOperation* forward_op_) {
+  forward_op_->op_name = op;
+  forward_op_->attrs.Reset(op);
+  return op_->Reset(op, raw_device_name);
+}
+absl::Status AddInput(AbstractOperation* op_, AbstractTensorHandle* input,
+                      ForwardOperation* forward_op_) {
+  TF_RETURN_IF_ERROR(op_->AddInput(input));
+  forward_op_->inputs.push_back(input);
+  return absl::OkStatus();
+}
+absl::Status AddInputList(AbstractOperation* op_,
+                          absl::Span<AbstractTensorHandle* const> inputs,
+                          ForwardOperation* forward_op_) {
+  TF_RETURN_IF_ERROR(op_->AddInputList(inputs));
+  for (auto input : inputs) {
+    forward_op_->inputs.push_back(input);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status SetAttrString(AbstractOperation* op_, const char* attr_name,
+                           const char* data, size_t length,
+                           ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name, absl::string_view(data, length));
+  return op_->SetAttrString(attr_name, data, length);
+}
+absl::Status SetAttrInt(AbstractOperation* op_, const char* attr_name,
+                        int64_t value, ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name, static_cast<int64_t>(value));
+  return op_->SetAttrInt(attr_name, value);
+}
+absl::Status SetAttrFloat(AbstractOperation* op_, const char* attr_name,
+                          float value, ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name, value);
+  return op_->SetAttrFloat(attr_name, value);
+}
+absl::Status SetAttrBool(AbstractOperation* op_, const char* attr_name,
+                         bool value, ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name, value);
+  return op_->SetAttrBool(attr_name, value);
+}
+absl::Status SetAttrType(AbstractOperation* op_, const char* attr_name,
+                         DataType value, ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name, value);
+  return op_->SetAttrType(attr_name, value);
+}
+absl::Status SetAttrShape(AbstractOperation* op_, const char* attr_name,
+                          const int64_t* dims, const int num_dims,
+                          ForwardOperation* forward_op_) {
+  if (num_dims > TensorShape::MaxDimensions()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Value specified for `", attr_name, "` has ", num_dims,
+                     " dimensions which is over the limit of ",
+                     TensorShape::MaxDimensions(), "."));
+  }
+  TensorShapeProto proto;
+  if (num_dims < 0) {
+    proto.set_unknown_rank(true);
+  } else {
+    for (int d = 0; d < num_dims; ++d) {
+      proto.add_dim()->set_size(dims[d]);
+    }
+  }
+
+  forward_op_->attrs.Set(attr_name, proto);
+  return op_->SetAttrShape(attr_name, dims, num_dims);
+}
+absl::Status SetAttrFunction(AbstractOperation* op_, const char* attr_name,
+                             const AbstractOperation* value,
+                             ForwardOperation* forward_op_) {
+  return absl::UnimplementedError(
+      "SetAttrFunction has not been implemented yet.");
+}
+absl::Status SetAttrFunctionName(AbstractOperation* op_, const char* attr_name,
+                                 const char* value, size_t length,
+                                 ForwardOperation* forward_op_) {
+  return absl::UnimplementedError(
+      "SetAttrFunctionName has not been implemented "
+      "yet.");
+}
+absl::Status SetAttrTensor(AbstractOperation* op_, const char* attr_name,
+                           AbstractTensorInterface* tensor,
+                           ForwardOperation* forward_op_) {
+  return absl::UnimplementedError(
+      "SetAttrTensor has not been implemented yet.");
+}
+absl::Status SetAttrStringList(AbstractOperation* op_, const char* attr_name,
+                               const void* const* values, const size_t* lengths,
+                               int num_values, ForwardOperation* forward_op_) {
+  std::vector<absl::string_view> v(num_values);
+  for (int i = 0; i < num_values; ++i) {
+    v[i] = absl::string_view(static_cast<const char*>(values[i]), lengths[i]);
+  }
+  forward_op_->attrs.Set(attr_name, v);
+  return op_->SetAttrStringList(attr_name, values, lengths, num_values);
+}
+absl::Status SetAttrFloatList(AbstractOperation* op_, const char* attr_name,
+                              const float* values, int num_values,
+                              ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name,
+                         gtl::ArraySlice<const float>(values, num_values));
+  return op_->SetAttrFloatList(attr_name, values, num_values);
+}
+absl::Status SetAttrIntList(AbstractOperation* op_, const char* attr_name,
+                            const int64_t* values, int num_values,
+                            ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(
+      attr_name, gtl::ArraySlice<const int64_t>(
+                     reinterpret_cast<const int64_t*>(values), num_values));
+  return op_->SetAttrIntList(attr_name, values, num_values);
+}
+absl::Status SetAttrTypeList(AbstractOperation* op_, const char* attr_name,
+                             const DataType* values, int num_values,
+                             ForwardOperation* forward_op_) {
+  forward_op_->attrs.Set(attr_name,
+                         gtl::ArraySlice<const DataType>(values, num_values));
+  return op_->SetAttrTypeList(attr_name, values, num_values);
+}
+absl::Status SetAttrBoolList(AbstractOperation* op_, const char* attr_name,
+                             const unsigned char* values, int num_values,
+                             ForwardOperation* forward_op_) {
+  std::unique_ptr<bool[]> b(new bool[num_values]);
+  for (int i = 0; i < num_values; ++i) {
+    b[i] = values[i];
+  }
+  forward_op_->attrs.Set(attr_name,
+                         gtl::ArraySlice<const bool>(b.get(), num_values));
+  return op_->SetAttrBoolList(attr_name, values, num_values);
+}
+absl::Status SetAttrShapeList(AbstractOperation* op_, const char* attr_name,
+                              const int64_t** dims, const int* num_dims,
+                              int num_values, ForwardOperation* forward_op_) {
+  std::unique_ptr<TensorShapeProto[]> proto(new TensorShapeProto[num_values]);
+  for (int i = 0; i < num_values; ++i) {
+    const auto num_dims_i = num_dims[i];
+
+    if (num_dims_i > TensorShape::MaxDimensions()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Value specified for `", attr_name, "` has ", num_dims_i,
+                       " dimensions which is over the limit of ",
+                       TensorShape::MaxDimensions(), "."));
+    }
+    if (num_dims_i < 0) {
+      proto[i].set_unknown_rank(true);
+    } else {
+      const int64_t* dims_i = dims[i];
+      auto proto_i = &proto[i];
+      for (int d = 0; d < num_dims_i; ++d) {
+        proto_i->add_dim()->set_size(dims_i[d]);
+      }
+    }
+  }
+  forward_op_->attrs.Set(
+      attr_name, absl::Span<const TensorShapeProto>(proto.get(), num_values));
+  return op_->SetAttrShapeList(attr_name, dims, num_dims, num_values);
+}
+absl::Status SetAttrFunctionList(AbstractOperation* op_, const char* attr_name,
+                                 absl::Span<const AbstractOperation*> values,
+                                 ForwardOperation* forward_op_) {
+  return absl::UnimplementedError(
+      "SetAttrFunctionList has not been "
+      "implemented yet.");
+}
+absl::Status Execute(AbstractOperation* op_, AbstractContext* ctx,
+                     absl::Span<AbstractTensorHandle*> retvals,
+                     int* num_retvals, ForwardOperation* forward_op_,
+                     Tape* tape, const GradientRegistry& registry) {
+  TF_RETURN_IF_ERROR(op_->Execute(retvals, num_retvals));
+  for (int i = 0; i < *num_retvals; i++) {
+    // TODO(srbs): Manage refcount of ForwardOperation's inputs/outputs.
+    forward_op_->outputs.push_back(retvals[i]);
+  }
+  // TODO(b/166669239): This is needed to support AttrBuilder::Get for string
+  // attributes. Number type attrs and DataType attrs work fine without this.
+  // Consider getting rid of this and making the behavior between number types
+  // and string consistent.
+  forward_op_->attrs.BuildNodeDef();
+  std::unique_ptr<GradientFunction> gradient_fn;
+  TF_RETURN_IF_ERROR(registry.Lookup(*forward_op_, &gradient_fn));
+  tape->RecordOperation(forward_op_->inputs, retvals.subspan(0, *num_retvals),
+                        gradient_fn.release(), op_->Name());
+  return absl::OkStatus();
+}
+}  // namespace internal
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/eager/gradients.h

@@ -0,0 +1,181 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_GRADIENTS_H_
+#define TENSORFLOW_C_EAGER_GRADIENTS_H_
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/tape.h"
+#include "tensorflow/core/common_runtime/eager/attr_builder.h"
+
+namespace tensorflow {
+namespace gradients {
+
+// =============== Experimental C++ API for computing gradients ===============
+
+// Sample gradient function:
+//
+// class AddGradientFunction : public GradientFunction {
+//  public:
+//   Status Compute(Context* ctx,
+//                  absl::Span<AbstractTensorHandle* const> grad_inputs,
+//                  absl::Span<AbstractTensorHandle*> grad_outputs) override {
+//     grad_outputs[0] = grad_inputs[0];
+//     grad_outputs[1] = grad_inputs[0];
+//     grad_outputs[0]->Ref();
+//     grad_outputs[1]->Ref();
+//     return OkStatus();
+//   }
+//   ~AddGradientFunction() override {}
+// };
+//
+// GradientFunction* AddRegisterer(const ForwardOperation& op) {
+//   // More complex gradient functions can use inputs/attrs etc. from the
+//   // forward `op`.
+//   return new AddGradientFunction;
+// }
+//
+// Status RegisterGradients(GradientRegistry* registry) {
+//   return registry->Register("Add", AddRegisterer);
+// }
+class GradientFunction {
+ public:
+  virtual absl::Status Compute(
+      AbstractContext* ctx,
+      absl::Span<AbstractTensorHandle* const> grad_outputs,
+      absl::Span<AbstractTensorHandle*> grad_inputs) = 0;
+  virtual ~GradientFunction() {}
+};
+
+// Metadata from the forward operation that is made available to the
+// gradient registerer to instantiate a GradientFunction.
+struct ForwardOperation {
+ public:
+  std::string op_name;
+  std::vector<AbstractTensorHandle*> inputs;
+  std::vector<AbstractTensorHandle*> outputs;
+  std::vector<int64_t> skip_input_indices;
+  AttrBuilder attrs;
+};
+
+using GradientFunctionFactory =
+    std::function<GradientFunction*(const ForwardOperation& op)>;
+
+// Map from op name to a `GradientFunctionFactory`.
+class GradientRegistry {
+ public:
+  absl::Status Register(const std::string& op,
+                        GradientFunctionFactory gradient_function_factory);
+  absl::Status Lookup(
+      const ForwardOperation& op,
+      std::unique_ptr<GradientFunction>* gradient_function) const;
+
+ private:
+  absl::flat_hash_map<std::string, GradientFunctionFactory> registry_;
+};
+
+// TODO(srbs): Figure out if we can avoid declaring this in the public header.
+// Wrapper for a tensor output of an operation executing under a tape.
+//
+// `GetID` returns a unique id for the wrapped tensor which is used to maintain
+// a map (`tensorflow::eager::TensorTape`) from the wrapped tensor to the id of
+// the op that produced it (or -1 if this tensor was watched using
+// `GradientTape::Watch`.) The op_id is simply a unique index assigned to each
+// op executed under the tape. A separate map (`tensorflow::eager::OpTape`)
+// maintains the map from `op_id` to a `OpTapeEntry` which stores the `op_type`,
+// inputs and outputs and the gradient function These data structures combined
+// allow us to trace the data dependencies between operations and hence compute
+// gradients.
+//
+// `ZerosLike` is not expected to be called and returns a nullptr. The creation
+// of default zeros grads is handled by the `DefaultGradientFunction` registered
+// for each op.
+// TODO(srbs): We need to define `ZerosLike` here to keep the compiler happy.
+// Figure out a way to avoid this.
+// TODO(srbs): Should ZerosLike check-fail instead of returning nullptr?
+class TapeTensor {
+ public:
+  explicit TapeTensor(AbstractTensorHandle* handle);
+  TapeTensor(const TapeTensor& other);
+  TapeTensor& operator=(const TapeTensor& other);
+  TapeTensor(TapeTensor&& other) noexcept;
+  TapeTensor& operator=(TapeTensor&& other) noexcept;
+  ~TapeTensor();
+
+  int64_t GetID() const;
+  tensorflow::DataType GetDType() const;
+
+  AbstractTensorHandle* ZerosLike() const;
+
+  AbstractTensorHandle* GetHandle() const;
+
+ private:
+  AbstractTensorHandle* handle_;
+};
+
+// A tracing/immediate-execution agnostic tape.
+//
+// Gradient functions defined for this tape must support handling null incoming
+// gradients.
+class Tape : protected eager::GradientTape<AbstractTensorHandle,
+                                           GradientFunction, TapeTensor> {
+ public:
+  using GradientTape<AbstractTensorHandle, GradientFunction,
+                     TapeTensor>::GradientTape;
+  // Returns whether the tape is persistent, i.e., whether the tape will hold
+  // onto its internal state after a call to `ComputeGradient`.
+  using GradientTape<AbstractTensorHandle, GradientFunction,
+                     TapeTensor>::IsPersistent;
+
+  // Adds this tensor to the list of watched tensors.
+  //
+  // This is a no-op if the tensor is already being watched either from an
+  // earlier call to `GradientTape::Watch` or being an output of an op with
+  // watched inputs.
+  void Watch(const AbstractTensorHandle*);
+  // Records an operation with given inputs and outputs
+  // on the tape and marks all its outputs as watched if at
+  // least one input of the op is watched and has a trainable dtype.
+  // op_name is optional and is used for debugging only.
+  void RecordOperation(absl::Span<AbstractTensorHandle* const> inputs,
+                       absl::Span<AbstractTensorHandle* const> outputs,
+                       GradientFunction* gradient_function,
+                       const std::string& op_name = "");
+  // Returns whether any tensor in a list of tensors is being watched and has
+  // a trainable dtype.
+  bool ShouldRecord(
+      absl::Span<const AbstractTensorHandle* const> tensors) const;
+  // Unwatches this tensor on the tape. Mainly used for cleanup when deleting
+  // eager tensors.
+  void DeleteTrace(const AbstractTensorHandle*);
+
+  // Consumes the internal state of the tape (so cannot be called more than
+  // once unless the tape is persistent) and produces the gradient of the target
+  // tensors with respect to the source tensors. The output gradients are used
+  // if not empty and not null. The result is populated with one tensor per
+  // target element.
+  absl::Status ComputeGradient(
+      AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> targets,
+      absl::Span<AbstractTensorHandle* const> sources,
+      absl::Span<AbstractTensorHandle* const> output_gradients,
+      absl::Span<AbstractTensorHandle*> result);
+};
+
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_GRADIENTS_H_

tensorflow/c/eager/gradients_internal.h

@@ -0,0 +1,93 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_GRADIENTS_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_GRADIENTS_INTERNAL_H_
+
+#include "tensorflow/c/eager/gradients.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+
+// Helper functions which delegate to `AbstractOperation`, update
+// the state of the ForwardOperation and call the tape as appropriate.
+// These APIs are mainly to facilitate testing and are subject to change.
+
+// Records the op name in the `ForwardOperation`.
+absl::Status Reset(AbstractOperation*, const char* op,
+                   const char* raw_device_name, ForwardOperation*);
+
+// Records the inputs in the `ForwardOperation`.
+absl::Status AddInput(AbstractOperation*, AbstractTensorHandle*,
+                      ForwardOperation*);
+absl::Status AddInputList(AbstractOperation*,
+                          absl::Span<AbstractTensorHandle* const> inputs,
+                          ForwardOperation*);
+
+// Sets the attrs in the `ForwardOperation`.
+absl::Status SetAttrString(AbstractOperation*, const char* attr_name,
+                           const char* data, size_t length, ForwardOperation*);
+absl::Status SetAttrInt(AbstractOperation*, const char* attr_name,
+                        int64_t value, ForwardOperation*);
+absl::Status SetAttrFloat(AbstractOperation*, const char* attr_name,
+                          float value, ForwardOperation*);
+absl::Status SetAttrBool(AbstractOperation*, const char* attr_name, bool value,
+                         ForwardOperation*);
+absl::Status SetAttrType(AbstractOperation*, const char* attr_name,
+                         DataType value, ForwardOperation*);
+absl::Status SetAttrShape(AbstractOperation*, const char* attr_name,
+                          const int64_t* dims, const int num_dims,
+                          ForwardOperation*);
+absl::Status SetAttrFunction(AbstractOperation*, const char* attr_name,
+                             const AbstractOperation* value, ForwardOperation*);
+absl::Status SetAttrFunctionName(AbstractOperation*, const char* attr_name,
+                                 const char* value, size_t length,
+                                 ForwardOperation*);
+absl::Status SetAttrTensor(AbstractOperation*, const char* attr_name,
+                           AbstractTensorInterface* tensor, ForwardOperation*);
+absl::Status SetAttrStringList(AbstractOperation*, const char* attr_name,
+                               const void* const* values, const size_t* lengths,
+                               int num_values, ForwardOperation*);
+absl::Status SetAttrFloatList(AbstractOperation*, const char* attr_name,
+                              const float* values, int num_values,
+                              ForwardOperation*);
+absl::Status SetAttrIntList(AbstractOperation*, const char* attr_name,
+                            const int64_t* values, int num_values,
+                            ForwardOperation*);
+absl::Status SetAttrTypeList(AbstractOperation*, const char* attr_name,
+                             const DataType* values, int num_values,
+                             ForwardOperation*);
+absl::Status SetAttrBoolList(AbstractOperation*, const char* attr_name,
+                             const unsigned char* values, int num_values,
+                             ForwardOperation*);
+absl::Status SetAttrShapeList(AbstractOperation*, const char* attr_name,
+                              const int64_t** dims, const int* num_dims,
+                              int num_values, ForwardOperation*);
+absl::Status SetAttrFunctionList(AbstractOperation*, const char* attr_name,
+                                 absl::Span<const AbstractOperation*> values,
+                                 ForwardOperation*);
+
+// Make the call to `Tape::RecordOperation`.
+absl::Status Execute(AbstractOperation*, AbstractContext*,
+                     absl::Span<AbstractTensorHandle*> retvals,
+                     int* num_retvals, ForwardOperation*, Tape*,
+                     const GradientRegistry&);
+
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_GRADIENTS_INTERNAL_H_

tensorflow/c/eager/gradients_test.cc

@@ -0,0 +1,241 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/gradients.h"
+
+#include <memory>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/gradients_internal.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/gradients/array_grad.h"
+#include "tensorflow/c/experimental/gradients/math_grad.h"
+#include "tensorflow/c/experimental/gradients/not_differentiable.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
+#include "tensorflow/c/experimental/ops/array_ops.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+using std::vector;
+using tensorflow::TF_StatusPtr;
+using tracing::TracingOperation;
+
+class CppGradients
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    absl::Status s = StatusFromTF_Status(status.get());
+    CHECK_EQ(errors::OK, s.code()) << s.message();
+  }
+};
+
+absl::Status RegisterGradients(GradientRegistry* registry) {
+  TF_RETURN_IF_ERROR(RegisterNotDifferentiable(registry, "CheckNumerics"));
+  return absl::OkStatus();
+}
+
+TEST_P(CppGradients, TestSetAttrString) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr t;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 1.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    t.reset(x_raw);
+  }
+
+  AbstractOperationPtr check_numerics_op(ctx->CreateOperation());
+  ForwardOperation forward_op;
+  absl::Status s = Reset(check_numerics_op.get(), "CheckNumerics",
+                         /*raw_device_name=*/nullptr, &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  if (isa<TracingOperation>(check_numerics_op.get())) {
+    s = dyn_cast<TracingOperation>(check_numerics_op.get())
+            ->SetOpName("check_numerics");
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+  }
+  s = AddInput(check_numerics_op.get(), t.get(), &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  std::string message = "This is the way!";
+  s = SetAttrString(check_numerics_op.get(), "message", message.data(),
+                    message.length(), &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  int num_retvals = 1;
+  std::vector<AbstractTensorHandle*> outputs(1);
+  GradientRegistry registry;
+  s = RegisterGradients(&registry);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  auto tape = std::make_unique<Tape>(/*persistent=*/false);
+  s = Execute(check_numerics_op.get(), ctx.get(), absl::MakeSpan(outputs),
+              &num_retvals, &forward_op, tape.get(), registry);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+
+  std::string read_message;
+  s = forward_op.attrs.Get("message", &read_message);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  ASSERT_EQ(read_message, message);
+}
+
+absl::Status RecordOperationWithNullGradientFunctionModel(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
+    absl::Span<AbstractTensorHandle*> outputs) {
+  Tape tape(/*persistent=*/false);
+  tape.Watch(inputs[0]);
+  AbstractTensorHandle* neg_output;
+  TF_RETURN_IF_ERROR(ops::Neg(ctx, inputs[0], &neg_output, "Neg"));
+  tape.RecordOperation(inputs, {neg_output}, nullptr, "Neg");
+  return tape.ComputeGradient(ctx,
+                              /*targets=*/{neg_output},
+                              /*sources=*/inputs,
+                              /*output_gradients=*/{}, outputs);
+}
+
+TEST_P(CppGradients, TestRecordOperationWithNullGradientFunctionRaises) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  std::vector<AbstractTensorHandle*> outputs(1);
+  absl::Status s = RunModel(RecordOperationWithNullGradientFunctionModel,
+                            ctx.get(), {x.get()}, absl::MakeSpan(outputs),
+                            /*use_function=*/!std::get<2>(GetParam()));
+  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
+  ASSERT_EQ(
+      "Provided null gradient_function for 'Neg'.\nIf the intent is to treat "
+      "this op as non-differentiable consider using RegisterNotDifferentiable "
+      "or NotDifferentiableGradientFunction.",
+      s.message());
+  ASSERT_EQ(nullptr, outputs[0]);
+}
+
+TEST_P(CppGradients, TestExecuteWithLargerOutputsVectorDoesNotCrash) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr t;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 1.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    t.reset(x_raw);
+  }
+
+  AbstractOperationPtr check_numerics_op(ctx->CreateOperation());
+  ForwardOperation forward_op;
+  absl::Status s = Reset(check_numerics_op.get(), "CheckNumerics",
+                         /*raw_device_name=*/nullptr, &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  if (isa<TracingOperation>(check_numerics_op.get())) {
+    s = dyn_cast<TracingOperation>(check_numerics_op.get())
+            ->SetOpName("check_numerics");
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+  }
+  s = AddInput(check_numerics_op.get(), t.get(), &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  std::string message = "This is the way!";
+  s = SetAttrString(check_numerics_op.get(), "message", message.data(),
+                    message.length(), &forward_op);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+
+  int num_retvals = 1;
+  // Allocate outputs with size 2, but we only expect 1 output.
+  // The second element will be initialized to nullptr.
+  std::vector<AbstractTensorHandle*> outputs(2, nullptr);
+
+  GradientRegistry registry;
+  s = RegisterGradients(&registry);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  auto tape = std::make_unique<Tape>(/*persistent=*/false);
+
+  // This call should NOT crash.
+  s = Execute(check_numerics_op.get(), ctx.get(), absl::MakeSpan(outputs),
+              &num_retvals, &forward_op, tape.get(), registry);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  EXPECT_EQ(num_retvals, 1);
+  EXPECT_NE(outputs[0], nullptr);
+  EXPECT_EQ(outputs[1], nullptr);  // Should remain nullptr
+}
+
+// TODO(b/164171226): Enable this test with tfrt after AddInputList is
+// supported. It is needed for IdentityN.
+#ifdef PLATFORM_GOOGLE
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*executing_eagerly*/ ::testing::Values(true, false)));
+#else
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*executing_eagerly*/ ::testing::Values(true, false)));
+#endif
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/eager/graph_function.cc

@@ -0,0 +1,36 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/graph_function.h"
+
+#include <utility>
+
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace tracing {
+namespace graph {
+GraphFunction::GraphFunction(FunctionDef fdef)
+    : AbstractFunction(kGraph),
+      func_record_(new FunctionRecord(std::move(fdef), {}, true)) {}
+GraphFunction::~GraphFunction() {}
+absl::Status GraphFunction::GetFunctionDef(const FunctionDef **fdef) {
+  *fdef = &(func_record_->fdef());
+  return absl::OkStatus();
+}
+}  // namespace graph
+}  // namespace tracing
+}  // namespace tensorflow

tensorflow/c/eager/graph_function.h

@@ -0,0 +1,53 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_GRAPH_FUNCTION_H_
+#define TENSORFLOW_C_EAGER_GRAPH_FUNCTION_H_
+
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/platform/refcount.h"
+namespace tensorflow {
+namespace tracing {
+namespace graph {
+using tensorflow::AbstractFunction;
+// Thin wrapper around a FunctionDef.
+class GraphFunction : public AbstractFunction {
+ public:
+  explicit GraphFunction(FunctionDef fdef);
+  ~GraphFunction() override;
+
+  // GraphFunction maybe stay alive for the duration of the returned
+  // FunctionDef.
+  absl::Status GetFunctionDef(const FunctionDef** fdef) override;
+
+  // Returns a shared reference to the wrapped function.
+  absl::StatusOr<core::RefCountPtr<FunctionRecord>> GetFunctionRecord()
+      override {
+    return func_record_.GetNewRef();
+  }
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractFunction* ptr) {
+    return ptr->getKind() == kGraph;
+  }
+
+ private:
+  core::RefCountPtr<FunctionRecord> func_record_;
+};
+}  // namespace graph
+}  // namespace tracing
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_GRAPH_FUNCTION_H_

tensorflow/c/eager/immediate_execution_context.h

@@ -0,0 +1,294 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <vector>
+
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/immediate_execution_distributed_manager.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/platform.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+class EagerExecutor;
+class EagerContext;
+class CustomDevice;
+class CustomDeviceOpHandler;
+class Device;
+
+// LINT.IfChange
+// Note: Keep in sync with exported copy of enum in eager/c_api.h.
+enum ContextDevicePlacementPolicy {
+  // Running operations with input tensors on the wrong device will fail.
+  DEVICE_PLACEMENT_EXPLICIT = 0,
+  // Copy the tensor to the right device but log a warning.
+  DEVICE_PLACEMENT_WARN = 1,
+  // Silently copy the tensor, which has a performance cost since the operation
+  // will be blocked till the copy completes. This is the default policy.
+  DEVICE_PLACEMENT_SILENT = 2,
+  // Placement policy which silently copies int32 tensors but not other dtypes.
+  DEVICE_PLACEMENT_SILENT_FOR_INT32 = 3,
+};
+// LINT.ThenChange(//tensorflow/c/eager/c_api.h)
+
+// Abstract interface to a context.
+//
+// A context is responsible for creating key objects such as Tensors,
+// TensorHandles & Operations.
+class ImmediateExecutionContext : public AbstractContext {
+ public:
+  // Optimized scalar creation functions
+  virtual AbstractTensorInterface* CreateInt64Scalar(int64_t value) = 0;
+  virtual AbstractTensorInterface* CreateUint64Scalar(uint64_t value) = 0;
+  virtual AbstractTensorInterface* CreateInt32Scalar(int32_t value) = 0;
+  virtual AbstractTensorInterface* CreateFloatScalar(float value) = 0;
+  virtual AbstractTensorInterface* CreateDoubleScalar(double value) = 0;
+  virtual AbstractTensorInterface* CreateHalfScalar(Eigen::half value) = 0;
+  virtual AbstractTensorInterface* CreateStringScalar(tstring value) = 0;
+  virtual AbstractTensorInterface* CreateComplex128Scalar(complex128 value) = 0;
+  virtual AbstractTensorInterface* CreateBoolScalar(bool value) = 0;
+
+  // Tensor creation functions
+  virtual AbstractTensorInterface* CreateTensor(
+      DataType dtype, absl::Span<const int64_t> dim_sizes) = 0;
+
+  typedef void (*MemoryReleaser)(void* data, size_t len, void* arg);
+
+  // Create a tensor instance from the given data buffer and description.
+  // `memory_releaser` will be called on destruction, and it's responsible for
+  // cleaning up the underlying buffer.
+  virtual AbstractTensorInterface* CreateTensor(
+      DataType dtype, const int64_t* dims, int num_dims, void* data, size_t len,
+      MemoryReleaser memory_releaser, void* memory_releaser_arg) = 0;
+
+  // Create a handle to wrap and manage a Tensor
+  virtual ImmediateExecutionTensorHandle* CreateLocalHandle(
+      AbstractTensorInterface* t) = 0;
+  // Copy the handle to another device.
+  virtual ImmediateExecutionTensorHandle* CopyTensorHandleToDevice(
+      ImmediateExecutionTensorHandle* handle, const char* device_name,
+      absl::Status* status) = 0;
+
+  // Create an operation to perform op execution
+  ImmediateExecutionOperation* CreateOperation() override = 0;
+
+  // Returns whether the runtime is backed by TFRT or the legacy TF Eager
+  // Runtime. This is necessary to decouple runtime-dependent
+  // code that is layered on top of the runtime.
+  virtual bool UsesTFRT() = 0;
+
+  // List attributes of available devices
+  virtual void ListDevices(std::vector<DeviceAttributes>* devices) = 0;
+
+  // Add `devices` into context's device manager. Context's device manager
+  // will take ownership and maintain devices' lifetime.
+  virtual absl::Status AddDevices(
+      std::vector<std::unique_ptr<Device>> devices) = 0;
+
+  // Block until all pending nodes are finished.
+  virtual absl::Status AsyncWait() = 0;
+
+  // Add a function (serialized FunctionDef protocol buffer) so that it can
+  // be executed as an op. Return error if the function with the same name
+  // already exists.
+  virtual absl::Status AddFunctionDef(const FunctionDef& fdef) = 0;
+
+  // Notifies about the function removal.
+  virtual absl::Status AddRemoveFunctionNotifier(
+      const std::string& func, std::function<void()> notifier) = 0;
+
+  // Same as `AddFunctionDef`, but additionally saves the `stack_traces` under
+  // the key of the function definition name (to be retrieved during function
+  // instantiation).
+  virtual absl::Status AddFunctionDefWithStackTraces(
+      const FunctionDef& fdef, const StackTracesMap& stack_traces) = 0;
+
+  // Find and return a added function by its name.
+  virtual const FunctionDef* FindFunctionDef(const std::string& name) const = 0;
+
+  // Find and return a function record added by its name.
+  virtual core::RefCountPtr<FunctionRecord> FindRecord(
+      const std::string& name) const = 0;
+
+  // Return the ParsedName of Host CPU device.
+  virtual const DeviceNameUtils::ParsedName& HostCPUParsedName() const = 0;
+  virtual const std::string& HostCPUName() const = 0;
+
+  // Configure soft device placement policy.
+  virtual void SetAllowSoftPlacement(bool enable) = 0;
+
+  // Configure device placement policy logging.
+  virtual void SetLogDevicePlacement(bool enable) = 0;
+
+  // Enables running eager ops as functions.
+  virtual void SetRunEagerOpAsFunction(bool enable) = 0;
+
+  // Enables rewriting jit_compile functions.
+  virtual void SetJitCompileRewrite(bool enable) = 0;
+
+  // Sets the device placement policy for the current thread.
+  virtual void SetThreadLocalDevicePlacementPolicy(
+      ContextDevicePlacementPolicy policy) = 0;
+  // Returns the device placement policy for the current thread.
+  virtual ContextDevicePlacementPolicy GetDevicePlacementPolicy() const = 0;
+
+  // Configure graph collection in RunMetadata.
+  virtual void SetShouldStoreGraphs(bool value) = 0;
+
+  // Return the collected RunMetadata. This method will transfer the ownership
+  // to the caller.
+  virtual std::unique_ptr<RunMetadata> ExportRunMetadata() = 0;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractContext* ptr) {
+    return ptr->getKind() == kEager || ptr->getKind() == kTfrt;
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Experimental Custom Device.
+  //===--------------------------------------------------------------------===//
+  virtual CustomDeviceOpHandler& GetCustomDeviceOpHandler() = 0;
+
+  // Returns whether `device_name` is registered as a custom device.
+  virtual bool IsCustomDevice(const std::string& device_name) = 0;
+
+  // Register a custom device. It will return error is the device name is
+  // already registered.
+  // TODO(tfrt-devs): Remove this method. Let caller register it directly into
+  // CustomDeviceOpHandler.
+  virtual absl::Status RegisterCustomDevice(
+      const std::string& name, std::unique_ptr<CustomDevice> device) = 0;
+
+  // Return FunctionLibraryDefinition. Transformations need to use it to use it
+  // to invoke MLIR compiler passes.
+  virtual FunctionLibraryDefinition* FuncLibDef() = 0;
+
+  // Resets the global rendezvous used for functions.
+  virtual void ResetGlobalRendezvousForFunction() = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Following are features in current TF Eager Runtime.
+  // TODO(tfrt-devs): Figure out a way to deprecate following features after
+  // migrated to TFRT.
+  //===--------------------------------------------------------------------===//
+  // Clear pending nodes in thread executors and kernel caches.
+  virtual void ClearCachesAndThreadExecutors() = 0;
+
+  // Initialize the step resource container for a training step. This is used
+  // in current TF runtime. For tfrt, it is used by fallback op handler.
+  virtual void StartStep() = 0;
+  // Destroy the step resource container for a training step.
+  virtual void EndStep() = 0;
+
+  // Return the Eager Executor for current thread. Please note that Eager
+  // Executor is only used in current TF but not in TFRT.
+  virtual EagerExecutor& Executor() = 0;
+  // Update the Eager Executor for current thread.
+  virtual void SetExecutorForThread(EagerExecutor* executor) = 0;
+
+  // Return a list of local tensorflow::Device*.
+  // TODO(tfrt-devs): We shouldn't expose legacy device in this API.
+  virtual std::vector<tensorflow::Device*> ListLocalTfDevices() = 0;
+
+  // Return a list of all tensorflow::Device*.
+  virtual std::vector<tensorflow::Device*> ListAllTfDevices() = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Following are helper functions to assist integrating TFRT with current
+  // TF eager runtime.
+  // TODO(b/172877902): These helper functions are currently used to support
+  // PyFuncOp on TFRT, and might be useful for ops that directly use low
+  // level TF APIs. Remove/replace the following functions when TFRT native
+  // ops are implemented.
+  //===--------------------------------------------------------------------===//
+  // Create an abstract tensor handle from tensorflow::Tensor.
+  virtual ImmediateExecutionTensorHandle* CreateLocalHandleFromTFTensor(
+      tensorflow::Tensor& t, const char* d_name) = 0;
+
+  // Convert a TFRT TensorHandle to tensorflow::TensorHandle.
+  virtual ImmediateExecutionTensorHandle* TFTensorHandleFromInterface(
+      ImmediateExecutionTensorHandle* handle) = 0;
+
+  virtual std::vector<std::string> GetLoggedOpsTestonly() { return {}; }
+
+  // Get a list of the names of functions that have been registered.
+  virtual std::vector<std::string> ListFunctionNames() = 0;
+
+  struct CacheStats {
+    int64_t kernel_cache_size;
+    int64_t device_cache_size;
+    std::map<std::string, int64_t> func_kernel_cache_entries;
+    int64_t local_rendezvous_cache_active_size;
+  };
+  virtual CacheStats GetCacheStats() = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Distributed runtime related functions.
+  //===--------------------------------------------------------------------===//
+#if !defined(IS_MOBILE_PLATFORM)
+  // Set up a multi-client distributed execution environment. Must be called on
+  // all tasks in the cluster.
+  // This call internally coordinates with other tasks to initialize the eager
+  // context and TF server for multi-client execution.
+  virtual absl::Status EnableCollectiveOps(const ServerDef& server_def) = 0;
+
+  // Set a distributed manager that helps set up, update, and check liveness
+  // of member tasks in the cluster.
+  virtual void SetDistributedManager(
+      std::unique_ptr<ImmediateExecutionDistributedManager> distributed) = 0;
+
+  virtual ImmediateExecutionDistributedManager* GetDistributedManager() = 0;
+#endif  // !IS_MOBILE_PLATFORM
+
+ protected:
+  explicit ImmediateExecutionContext(AbstractContextKind kind)
+      : AbstractContext(kind) {}
+  ~ImmediateExecutionContext() override {}
+};
+
+namespace internal {
+struct ImmediateExecutionContextDeleter {
+  void operator()(ImmediateExecutionContext* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+}  // namespace internal
+
+using ImmediateContextPtr =
+    std::unique_ptr<ImmediateExecutionContext,
+                    internal::ImmediateExecutionContextDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_

tensorflow/c/eager/immediate_execution_distributed_manager.h

@@ -0,0 +1,70 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_DISTRIBUTED_MANAGER_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_DISTRIBUTED_MANAGER_H_
+
+#include <cstdint>
+#include <string>
+
+#include "tensorflow/core/platform/status.h"
+
+namespace tsl {
+class CoordinationServiceAgent;
+}
+
+namespace tensorflow {
+class ImmediateExecutionContext;
+class ServerDef;
+class WorkerEnv;
+class WorkerCacheInterface;
+
+class ImmediateExecutionDistributedManager {
+ public:
+  virtual ~ImmediateExecutionDistributedManager() {}
+
+  // Set up distributed execution environment on local and remote tasks.
+  // When `reset_context` is true, initialize new cluster context state based
+  // on cluster configurations provided in `server_def`; otherwise, update
+  // existing context state with the provided `server_def`. Contexts created
+  // on remote tasks will be considered stale and garbage collected after
+  // `keep_alive_secs` of inactivity.
+  virtual absl::Status SetOrUpdateServerDef(
+      const ServerDef& server_def, bool reset_context, int keep_alive_secs,
+      int64_t init_timeout_in_ms, int retries,
+      bool clear_existing_contexts = false) = 0;
+
+  // Initializes context for the local worker and no contexts will be created
+  // for remote workers. Currently this only works for resetting context.
+  // TODO(b/289445025): Consider removing this when we find a proper fix.
+  virtual absl::Status InitializeLocalOnlyContext(const ServerDef& server_def,
+                                                  int keep_alive_secs) = 0;
+
+  // Set up a multi-client distributed execution environment. Must be called
+  // on all tasks in the cluster. This call internally coordinates with other
+  // tasks to initialize the eager context and TF server for multi-client
+  // execution.
+  virtual absl::Status EnableCollectiveOps(const ServerDef& server_def) = 0;
+
+  // Check if the remote task is alive.
+  virtual absl::Status CheckRemoteAlive(const std::string& remote_task_name,
+                                        bool* is_alive) = 0;
+
+  // Get pointer to the coordination service agent instance.
+  virtual tsl::CoordinationServiceAgent* GetCoordinationServiceAgent() = 0;
+};
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_DISTRIBUTED_MANAGER_H_

tensorflow/c/eager/immediate_execution_operation.h

@@ -0,0 +1,108 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "absl/status/status.h"
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/managed_stack_trace.h"
+
+struct TFE_Op;
+
+namespace tensorflow {
+
+class ImmediateExecutionContext;
+class AbstractOpAttrs;
+
+// Abstract interface to an operation.
+class ImmediateExecutionOperation : public AbstractOperation {
+ public:
+  virtual void Clear() = 0;
+
+  // Returns the inputs of this op.
+  virtual absl::Span<ImmediateExecutionTensorHandle* const> GetInputs()
+      const = 0;
+  virtual absl::Status SetInput(size_t index,
+                                ImmediateExecutionTensorHandle* input) = 0;
+
+  virtual ImmediateExecutionContext* GetContext() const = 0;
+
+  // Following two methods are used to support custom device.
+  // Return true if the inputs contain custom device tensor handle. It means
+  // that the argument need to be handled by a custom device.
+  virtual bool HasCustomDeviceInput() const = 0;
+
+  virtual const tensorflow::OpDef* OpDef() const = 0;
+
+  virtual absl::Status InputLength(const char* input_name, int* length) = 0;
+  virtual absl::Status OutputLength(const char* output_name, int* length) = 0;
+
+  // Set stack trace to be used for potential async error reporting.
+  virtual void SetStackTrace(ManagedStackTrace stack_trace) = 0;
+
+  virtual const tensorflow::AbstractOpAttrs* GetOpAttrs() const = 0;
+  virtual void AddAttrs(const AbstractOpAttrs* op_attrs) = 0;
+
+  virtual absl::Status SetAttrValue(const char* attr_name,
+                                    const AttrValue& value) = 0;
+
+  virtual void SetCancellationManager(
+      CancellationManager* cancellation_manager) = 0;
+
+  // Returns the stack trace set by `SetStackTrace` if exists.
+  virtual absl::optional<ManagedStackTrace> GetStackTrace() = 0;
+
+  virtual void SetStepId(int64_t step_id) = 0;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractOperation* ptr) {
+    return ptr->getKind() == kEager || ptr->getKind() == kTfrt;
+  }
+
+ protected:
+  explicit ImmediateExecutionOperation(AbstractOperationKind kind)
+      : AbstractOperation(kind) {}
+  ~ImmediateExecutionOperation() override = default;
+};
+
+namespace internal {
+struct ImmediateExecutionOperationDeleter {
+  void operator()(ImmediateExecutionOperation* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+}  // namespace internal
+
+using ImmediateOpPtr =
+    std::unique_ptr<ImmediateExecutionOperation,
+                    internal::ImmediateExecutionOperationDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_

tensorflow/c/eager/immediate_execution_tensor_handle.cc

@@ -0,0 +1,61 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+
+namespace tensorflow {
+
+std::string ImmediateExecutionTensorHandle::DebugString() const {
+  PartialTensorShape shape;
+  std::string shape_string;
+  if (Shape(&shape).ok()) {
+    shape_string = shape.DebugString();
+  } else {
+    shape_string = "<error computing shape>";
+  }
+  std::string value_string;
+  if (!SummarizeValue(value_string).ok()) {
+    value_string = "<error computing value>";
+  }
+  if (value_string.length() > 100) {
+    // The default NumPy-style output can be distractingly long in error
+    // messages.
+    value_string = absl::StrCat(value_string.substr(0, 100), " [...]");
+  }
+  absl::Status s;
+  const char* device_name = DeviceName(&s);
+  if (!s.ok()) {
+    device_name = "<error fetching device name>";
+  }
+  return absl::StrCat("TensorHandle(", value_string, ", shape=", shape_string,
+                      ", dtype=", DataType_Name(DataType()), ", device=\"",
+                      device_name, "\")");
+}
+
+absl::Status ImmediateExecutionTensorHandle::SummarizeValue(
+    std::string& summary) const {
+  absl::Status status;
+  AbstractTensorPtr resolved(
+      // TODO(allenl): Resolve should be const, and the caches that get updated
+      // marked mutable.
+      const_cast<ImmediateExecutionTensorHandle*>(this)->Resolve(&status));
+  if (!status.ok()) {
+    return status;
+  }
+  summary = resolved->SummarizeValue();
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/c/eager/immediate_execution_tensor_handle.h

@@ -0,0 +1,105 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_
+
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Abstract interface to a TensorHandle.
+//
+// A TensorHandle is management class around a Tensor which may track additional
+// metadata and synchronization.
+//
+// This allows us to hide concrete implementations of TensorHandle from header
+// files. The interface lists the common functionality that must be provided by
+// any concrete implementation. However, in cases where the true concrete class
+// is needed a static_cast can be applied.
+class ImmediateExecutionTensorHandle : public AbstractTensorHandle {
+ public:
+  // Returns number of dimensions.
+  virtual absl::Status NumDims(int* num_dims) const = 0;
+  // Returns number of elements across all dimensions.
+  virtual absl::Status NumElements(int64_t* num_elements) const = 0;
+  // Returns size of specified dimension
+  //
+  // -1 indicates an unknown axis length; this is unreachable for most standard
+  // ImmediateExecutionTensorHandles, but comes up for example when computing
+  // the shape of a parallel tensor with component shapes differing across
+  // devices.
+  virtual absl::Status Dim(int dim_index, int64_t* dim) const = 0;
+
+  // Returns the device which created the handle.
+  virtual const char* DeviceName(absl::Status* status) const = 0;
+  // Returns the device where the tensor was placed.
+  virtual const char* BackingDeviceName(absl::Status* status) const = 0;
+  // Returns the device type which created the handle.
+  virtual const char* DeviceType(absl::Status* status) const = 0;
+  // Returns the device ID which created the handle.
+  virtual int DeviceId(absl::Status* status) const = 0;
+  // Returns a tensor for the handle. If tensor is remote, it will be copied.
+  virtual AbstractTensorInterface* Resolve(absl::Status* status) = 0;
+
+  std::string DebugString() const override;
+
+  // Returns a Boolean hint indicating whether callers should prefer
+  // `SummarizeValue` to resolving this handle and formatting the tensor.
+  //
+  // For example some tensor handles may represent distributed values, in which
+  // case placement information is lost when resolving the handle.
+  //
+  // If false, a caller might implement pretty-printing by resolving and
+  // iterating over the resulting tensor. This may still be viable if resolving
+  // the handle loses information, but `SummarizeValue` would be more precise.
+  virtual bool PreferCustomSummarizer() const { return false; }
+
+  // Returns a string which summarizes the value of this TensorHandle, for
+  // debugging. Does not include a shape or dtype.
+  //
+  // Included in the default implementation of DebugString.
+  virtual absl::Status SummarizeValue(std::string& summary) const;
+
+  // For LLVM style RTTI.
+  static bool classof(const AbstractTensorHandle* ptr) {
+    return ptr->getKind() == kEager || ptr->getKind() == kTfrt;
+  }
+
+ protected:
+  explicit ImmediateExecutionTensorHandle(AbstractTensorHandleKind kind)
+      : AbstractTensorHandle(kind) {}
+  ~ImmediateExecutionTensorHandle() override {}
+};
+
+namespace internal {
+struct ImmediateExecutionTensorHandleDeleter {
+  void operator()(ImmediateExecutionTensorHandle* p) const {
+    if (p != nullptr) {
+      p->Unref();
+    }
+  }
+};
+}  // namespace internal
+
+using ImmediateTensorHandlePtr =
+    std::unique_ptr<ImmediateExecutionTensorHandle,
+                    internal::ImmediateExecutionTensorHandleDeleter>;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_

tensorflow/c/eager/parallel_device/BUILD

@@ -0,0 +1,184 @@
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+# Currently pybind extension shared objects must use only C API headers since
+# the C API has static initializers duplicated in the Python bindings. So we
+# need a second rule that omits .cc files, in
+# tensorflow/python:_pywrap_parallel_device.
+filegroup(
+    name = "lib_headers",
+    srcs = ["parallel_device_lib.h"],
+)
+
+filegroup(
+    name = "lib_sources",
+    srcs = ["parallel_device_lib.cc"],
+)
+
+filegroup(
+    name = "device_headers",
+    srcs = ["parallel_device.h"],
+)
+
+filegroup(
+    name = "device_sources",
+    srcs = ["parallel_device.cc"],
+)
+
+filegroup(
+    name = "headers",
+    srcs = [
+        ":device_headers",
+        ":lib_headers",
+    ],
+    visibility = ["//tensorflow/python:__pkg__"],
+)
+
+filegroup(
+    name = "sources",
+    srcs = [
+        ":device_sources",
+        ":lib_sources",
+    ],
+    visibility = ["//tensorflow/python:__pkg__"],
+)
+
+cc_library(
+    name = "parallel_device",
+    srcs = [":device_sources"],
+    hdrs = [":device_headers"],
+    visibility = ["//visibility:private"],
+    deps = [
+        ":parallel_device_lib",
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:variant",
+    ],
+)
+
+cc_library(
+    name = "parallel_device_lib",
+    srcs = [":lib_sources"],
+    hdrs = [":lib_headers"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:safe_ptr",
+        "//tensorflow/c:tf_status_internal",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/c/eager:tfe_cancellation_manager_internal",
+        "//tensorflow/c/eager:tfe_op_internal",
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
+    ],
+)
+
+tf_cc_test(
+    name = "parallel_device_lib_test",
+    srcs = ["parallel_device_lib_test.cc"],
+    deps = [
+        ":parallel_device_lib",
+        ":parallel_device_testlib",
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:c_api_experimental",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:tfe_context_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "parallel_device_testlib",
+    testonly = 1,
+    srcs = ["parallel_device_testlib.cc"],
+    hdrs = ["parallel_device_testlib.h"],
+    deps = [
+        ":parallel_device",
+        ":parallel_device_lib",
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:c_api_experimental",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_test(
+    name = "parallel_device_test",
+    srcs = ["parallel_device_test.cc"],
+    deps = [
+        ":parallel_device_lib",
+        ":parallel_device_testlib",
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:c_api_experimental",
+        "//tensorflow/c:tf_status_internal",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tf_cc_test(
+    name = "parallel_device_remote_test",
+    srcs = ["parallel_device_remote_test.cc"],
+    # TODO(b/136478427): Enable global heap checking when servers shut down
+    # cleanly.
+    args = ["--heap_check="],
+    deps = [
+        ":parallel_device_lib",
+        ":parallel_device_testlib",
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:c_api_experimental",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/distributed_runtime:master_env",
+        "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib",
+        "//tensorflow/core/platform:strcat",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+    ],
+)

tensorflow/c/eager/parallel_device/parallel_device.cc

@@ -0,0 +1,441 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/parallel_device/parallel_device.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/variant.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_buffer.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+
+namespace tensorflow {
+namespace parallel_device {
+namespace {
+
+class OpDeleter {
+ public:
+  void operator()(TFE_Op* to_delete) const { TFE_DeleteOp(to_delete); }
+};
+
+using OpPtr = std::unique_ptr<TFE_Op, OpDeleter>;
+
+using MaybeParallelTensorOwned =
+    std::variant<std::unique_ptr<ParallelTensor>, TensorHandlePtr>;
+
+using MaybeParallelTensorUnowned =
+    std::variant<ParallelTensor*, TFE_TensorHandle*>;
+
+// A ParallelDevice on its own is not registered with a TFE_Context, and so has
+// no device name (e.g. for `tf.device`). `NamedParallelDevice` associates a
+// name with it, which lets us pack its `ParallelTensor`s into TFE_TensorHandles
+// placed on the parallel device.
+class NamedParallelDevice {
+ public:
+  NamedParallelDevice(const std::string& name,
+                      std::unique_ptr<ParallelDevice> parallel_device)
+      : device_name_(name), parallel_device_(std::move(parallel_device)) {}
+  const std::string& name() const { return device_name_; }
+  const ParallelDevice& device() const { return *parallel_device_; }
+
+ private:
+  std::string device_name_;
+  std::unique_ptr<ParallelDevice> parallel_device_;
+};
+
+std::optional<std::vector<MaybeParallelTensorOwned>> ExecuteWithSpecialOps(
+    const ParallelDevice& parallel_device,
+    const std::string& parallel_device_name, TFE_Context* context,
+    std::vector<MaybeParallelTensorUnowned> inputs, const char* operation_name,
+    const TFE_OpAttrs* attributes, int expected_max_outputs,
+    TF_Status* status) {
+  std::optional<std::vector<MaybeParallelTensorOwned>> result;
+  // TODO(allenl): We should remove "TPU" from these op names at the very least,
+  // or consider other ways of packing/unpacking parallel tensors.
+  if (operation_name == std::string("TPUReplicatedInput")) {
+    // Special-cased operation for packing per-device tensors into one parallel
+    // tensor.
+    if (inputs.size() != parallel_device.num_underlying_devices()) {
+      std::string message(absl::StrCat(
+          "The parallel device ", parallel_device_name, " expected ",
+          parallel_device.num_underlying_devices(),
+          " inputs to TPUReplicatedInput, but got ", inputs.size()));
+      TF_SetStatus(status, TF_INVALID_ARGUMENT, message.c_str());
+      return result;
+    }
+    std::vector<TensorHandlePtr> components;
+    components.reserve(inputs.size());
+    for (int i = 0; i < inputs.size(); ++i) {
+      if (absl::holds_alternative<ParallelTensor*>(inputs[i])) {
+        std::string message(absl::StrCat(
+            "Expected all inputs to TPUReplicatedInput to be non-parallel "
+            "TensorHandles. The input ",
+            i,
+            " was a parallel tensor (already "
+            "placed on the parallel device)."));
+        TF_SetStatus(status, TF_INVALID_ARGUMENT, message.c_str());
+        return result;
+      }
+      components.emplace_back(TFE_TensorHandleCopySharingTensor(
+          absl::get<TFE_TensorHandle*>(inputs[i]), status));
+    }
+    std::vector<MaybeParallelTensorOwned> result_content;
+    result_content.reserve(1);
+    result_content.push_back(ParallelTensor::FromTensorHandles(
+        parallel_device, std::move(components), status));
+    if (TF_GetCode(status) != TF_OK) return result;
+    result.emplace(std::move(result_content));
+    return result;
+  } else if (operation_name == std::string("TPUReplicatedOutput")) {
+    // Special-cased operation for un-packing one parallel tensor into
+    // per-device tensors.
+    OpPtr op(TFE_NewOp(context, operation_name, status));
+    TFE_OpAddAttrs(op.get(), attributes);
+    int expected_outputs = TFE_OpGetOutputLength(op.get(), "outputs", status);
+    if (TF_GetCode(status) != TF_OK) return result;
+    if (expected_outputs != parallel_device.num_underlying_devices()) {
+      std::string message(absl::StrCat(
+          "The parallel device ", parallel_device_name, " expected ",
+          parallel_device.num_underlying_devices(),
+          " outputs for TPUReplicatedOutput, but got ", expected_outputs));
+      TF_SetStatus(status, TF_INVALID_ARGUMENT, message.c_str());
+      return result;
+    }
+    if (inputs.size() != 1) {
+      std::string message(absl::StrCat("The parallel device ",
+                                       parallel_device_name,
+                                       " expected 1 input for "
+                                       "TPUReplicatedOutput, but got ",
+                                       inputs.size()));
+      TF_SetStatus(status, TF_INVALID_ARGUMENT, message.c_str());
+      return result;
+    }
+    if (absl::holds_alternative<TFE_TensorHandle*>(inputs[0])) {
+      TF_SetStatus(status, TF_INVALID_ARGUMENT,
+                   "Expected the input to "
+                   "TPUReplicatedOutput to be a parallel tensor (placed on the "
+                   "parallel device).");
+      return result;
+    }
+    ParallelTensor* t = absl::get<ParallelTensor*>(inputs[0]);
+    std::vector<MaybeParallelTensorOwned> outputs;
+    outputs.reserve(t->num_tensors());
+    for (int i = 0; i < t->num_tensors(); ++i) {
+      TensorHandlePtr this_output(
+          TFE_TensorHandleCopySharingTensor(t->tensor(i), status));
+      outputs.emplace_back(std::move(this_output));
+      if (TF_GetCode(status) != TF_OK) return result;
+    }
+    result.emplace(std::move(outputs));
+    return result;
+  }
+  std::vector<ParallelTensor*> parallel_inputs;
+  std::vector<std::unique_ptr<ParallelTensor>> implicitly_broadcast_tensors;
+  parallel_inputs.reserve(inputs.size());
+  implicitly_broadcast_tensors.reserve(inputs.size());  // not tight
+  for (const auto& input : inputs) {
+    if (absl::holds_alternative<TFE_TensorHandle*>(input)) {
+      if (operation_name == std::string("_EagerConst")) {
+        // Non-parallel tensors from _EagerConst/tf.constant are implicitly
+        // 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
+        // just copy-off but includes a sum) and consideration of performance.
+        //
+        // TODO(allenl): There may be smarter ways to do this copy in some
+        // cases, i.e. with a collective broadcast. We'll need to be careful
+        // about things that are taken as inputs on the host or on their
+        // existing device (for multi-device functions).
+        std::unique_ptr<ParallelTensor> parallel_tensor(
+            parallel_device.CopyToParallelDevice(
+                context, absl::get<TFE_TensorHandle*>(input), status));
+        if (TF_GetCode(status) != TF_OK) return std::nullopt;
+        parallel_inputs.push_back(parallel_tensor.get());
+        implicitly_broadcast_tensors.emplace_back(std::move(parallel_tensor));
+      } else {
+        TF_SetStatus(
+            status, TF_INVALID_ARGUMENT,
+            absl::StrCat(
+                "Got a non-parallel tensor ",
+                tensorflow::unwrap(absl::get<TFE_TensorHandle*>(input))
+                    ->DebugString(),
+                " as input to a parallel operation. First pack non-parallel "
+                "tensors for each device into a parallel tensor explicitly.")
+                .c_str());
+        return std::nullopt;
+      }
+    } else {
+      parallel_inputs.push_back(absl::get<ParallelTensor*>(input));
+    }
+  }
+  std::optional<std::vector<std::unique_ptr<ParallelTensor>>>
+      maybe_parallel_results(
+          parallel_device.Execute(context, parallel_inputs, operation_name,
+                                  attributes, expected_max_outputs, status));
+  if (!maybe_parallel_results.has_value()) return result;
+  std::vector<std::unique_ptr<ParallelTensor>> parallel_results(
+      std::move(maybe_parallel_results.value()));
+  std::vector<MaybeParallelTensorOwned> result_content;
+  result_content.reserve(parallel_results.size());
+  for (std::unique_ptr<ParallelTensor>& parallel_result : parallel_results) {
+    result_content.push_back(
+        MaybeParallelTensorOwned(std::move(parallel_result)));
+  }
+  result.emplace(std::move(result_content));
+  return result;
+}
+
+// Used as an argument to TFE_NewCustomDeviceTensorHandle, indicating how
+// ParallelTensors wrapped in TFE_TensorHandles should be cleaned up once their
+// reference counts drop to zero.
+void ParallelTensorDeallocator(void* data) {
+  delete reinterpret_cast<ParallelTensor*>(data);
+}
+
+// Used as an argument to TFE_NewCustomDeviceTensorHandle, for computing the
+// number of dimensions of a parallel tensor.
+int ParallelTensorNumDims(void* data, TF_Status* status) {
+  const std::vector<int64_t>* shape;
+  absl::Status s = reinterpret_cast<ParallelTensor*>(data)->Shape(&shape);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return -1;
+  }
+  return shape->size();
+}
+
+// Used as an argument to TFE_NewCustomDeviceTensorHandle, for computing a
+// dimension of a parallel tensor.
+int64_t ParallelTensorDim(void* data, int dim_index, TF_Status* status) {
+  const std::vector<int64_t>* shape;
+  absl::Status s = reinterpret_cast<ParallelTensor*>(data)->Shape(&shape);
+  if (!s.ok()) {
+    tsl::Set_TF_Status_from_Status(status, s);
+    return -1;
+  }
+  return (*shape)[dim_index];
+}
+
+TF_Buffer* ParallelTensorSummarize(void* data, TF_Status* status) {
+  ParallelTensor* parallel_tensor = reinterpret_cast<ParallelTensor*>(data);
+  std::string summary;
+  absl::Status cpp_status = parallel_tensor->SummarizeValue(summary);
+  if (!cpp_status.ok()) {
+    tsl::Set_TF_Status_from_Status(status, cpp_status);
+    return nullptr;
+  }
+  return TF_NewBufferFromString(summary.data(), summary.size());
+}
+
+TensorHandlePtr ParallelTensorToTensorHandle(
+    const std::string& parallel_device_name, TFE_Context* context,
+    std::unique_ptr<ParallelTensor> t, TF_Status* status) {
+  // The resulting TensorHandle owns an opaque pointer to "device memory", which
+  // for a ParallelDevice is really a ParallelTensor. When the TensorHandle is
+  // deleted, it will call ParallelTensorDeallocator to free the struct.
+  ParallelTensor* t_released = t.release();
+  TFE_CustomDeviceTensorHandleMethods handle_methods;
+  handle_methods.num_dims = &ParallelTensorNumDims;
+  handle_methods.dim = &ParallelTensorDim;
+  handle_methods.deallocator = &ParallelTensorDeallocator;
+  handle_methods.summarize = &ParallelTensorSummarize;
+  return TensorHandlePtr(TFE_NewCustomDeviceTensorHandle(
+      context, parallel_device_name.c_str(), t_released->dtype(), t_released,
+      handle_methods, status));
+}
+
+// For TFE_CustomDevice::copy_tensor_to_device in the parallel device
+// registration.
+//
+// Since this function is used to satisfy the TFE_CustomDevice C API,
+// device_info is passed in using a C-style generic. It must always be a
+// ParallelDevice.
+TFE_TensorHandle* CopyToParallelDevice(TFE_Context* context,
+                                       TFE_TensorHandle* tensor,
+                                       TF_Status* status, void* device_info) {
+  TF_SetStatus(
+      status, TF_UNIMPLEMENTED,
+      absl::StrCat("Trying to copy a tensor ",
+                   tensorflow::unwrap(tensor)->DebugString(),
+                   " on to a parallel device. Pack non-parallel "
+                   "tensors for each device into a parallel tensor explicitly.")
+          .c_str());
+  return nullptr;
+}
+
+// For TFE_CustomDevice::copy_tensor_from_device in the parallel device
+// registration.
+//
+// Currently this is an error, and un-packing ParallelTensors must be performed
+// explicitly by running a TPUReplicatedOutput operation on the parallel device.
+//
+// TODO(allenl): There are some use-cases that are only supported by copying to
+// host at the moment (e.g. debug print on a tensor, .numpy(), etc.). We either
+// need to return something here or address these use-cases one by one.
+TFE_TensorHandle* CopyTensorFromParallelDevice(TFE_Context* context,
+                                               TFE_TensorHandle* tensor,
+                                               const char* target_device_name,
+                                               TF_Status* status,
+                                               void* device_info) {
+  ParallelTensor* parallel_tensor = reinterpret_cast<ParallelTensor*>(
+      TFE_TensorHandleDevicePointer(tensor, status));
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  if (parallel_tensor->num_tensors() == 1) {
+    // Copy-off for single-device tensors is allowed to make debugging dynamic
+    // control flow easier.
+    return TFE_TensorHandleCopySharingTensor(parallel_tensor->tensor(0),
+                                             status);
+  } else {
+    TF_SetStatus(
+        status, TF_UNIMPLEMENTED,
+        absl::StrCat(
+            "Trying to copy a tensor out of a parallel device. Since there "
+            "are multiple components to parallel tensors, they must be "
+            "unpacked explicitly.\n",
+            tensorflow::unwrap(tensor)->DebugString())
+            .c_str());
+    return nullptr;
+  }
+}
+
+// For TFE_CustomDevice::execute in the parallel device registration.
+//
+// Since this function is used to satisfy the TFE_CustomDevice C API,
+// device_info is passed in using a C-style generic. It must always be a
+// ParallelDevice.
+void ParallelDeviceExecute(const TFE_Op* original_op, int* num_outputs,
+                           TFE_TensorHandle** outputs, TF_Status* status,
+                           void* device_info) {
+  const char* requested_placement = TFE_OpGetDevice(original_op, status);
+  if (*requested_placement == '\0') {
+    TF_SetStatus(
+        status, TF_INTERNAL,
+        "Ops must be placed on the parallel device explicitly, or their inputs "
+        "first un-packed. Got an un-placed op with an input placed on the "
+        "parallel device.");
+    return;
+  }
+  TFE_Context* context = TFE_OpGetContext(original_op, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  const char* operation_name = TFE_OpGetName(original_op, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  const TFE_OpAttrs* attributes = TFE_OpGetAttrs(original_op);
+
+  NamedParallelDevice* named_device =
+      reinterpret_cast<NamedParallelDevice*>(device_info);
+  std::vector<MaybeParallelTensorUnowned> typed_inputs;
+  int num_inputs = TFE_OpGetFlatInputCount(original_op, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  typed_inputs.reserve(num_inputs);
+  for (int i = 0; i < num_inputs; ++i) {
+    TFE_TensorHandle* input = TFE_OpGetFlatInput(original_op, i, status);
+    if (TF_GetCode(status) != TF_OK) return;
+    const char* tensor_handle_device =
+        TFE_TensorHandleDeviceName(input, status);
+    if (TF_GetCode(status) != TF_OK) return;
+    if (named_device->name() == tensor_handle_device) {
+      // We assume that any tensors already placed on this device are
+      // ParallelTensors.
+      typed_inputs.emplace_back(reinterpret_cast<ParallelTensor*>(
+          TFE_TensorHandleDevicePointer(input, status)));
+      if (TF_GetCode(status) != TF_OK) return;
+    } else {
+      typed_inputs.emplace_back(input);
+    }
+  }
+
+  std::optional<std::vector<MaybeParallelTensorOwned>> maybe_typed_outputs(
+      ExecuteWithSpecialOps(named_device->device(), named_device->name(),
+                            context, std::move(typed_inputs), operation_name,
+                            attributes, *num_outputs, status));
+  if (TF_GetCode(status) != TF_OK) return;
+  if (!maybe_typed_outputs.has_value()) {
+    TF_SetStatus(status, TF_INTERNAL, "OK status but no value was returned.");
+    return;
+  }
+
+  std::vector<MaybeParallelTensorOwned> typed_outputs(
+      std::move(maybe_typed_outputs.value()));
+
+  if (typed_outputs.size() > *num_outputs) {
+    TF_SetStatus(status, TF_INTERNAL,
+                 "The allocated output buffer was too small.");
+    return;
+  }
+
+  for (int i = 0; i < typed_outputs.size(); ++i) {
+    MaybeParallelTensorOwned typed_output(std::move(typed_outputs[i]));
+    if (absl::holds_alternative<TensorHandlePtr>(typed_output)) {
+      outputs[i] = absl::get<TensorHandlePtr>(typed_output).release();
+    } else {
+      outputs[i] = ParallelTensorToTensorHandle(
+                       named_device->name(), context,
+                       std::move(absl::get<std::unique_ptr<ParallelTensor>>(
+                           typed_output)),
+                       status)
+                       .release();
+      if (TF_GetCode(status) != TF_OK) return;
+    }
+  }
+  *num_outputs = typed_outputs.size();
+}
+
+// For TFE_CustomDevice::delete_device in the parallel device registration.
+//
+// Since this function is used to satisfy the TFE_CustomDevice C API,
+// device_info is passed in using a C-style generic. It must always be a
+// ParallelDevice.
+void DeleteParallelDevice(void* device_info) {
+  delete reinterpret_cast<NamedParallelDevice*>(device_info);
+}
+
+}  // namespace
+
+void AllocateParallelDevice(const char* device_name,
+                            const char* const* underlying_devices,
+                            int num_underlying_devices,
+                            TFE_CustomDevice* device, void** device_info) {
+  device->copy_tensor_to_device = &CopyToParallelDevice;
+  device->copy_tensor_from_device = &CopyTensorFromParallelDevice;
+  device->delete_device = &DeleteParallelDevice;
+  device->execute = &ParallelDeviceExecute;
+  std::vector<std::string> underlying_devices_vector;
+  underlying_devices_vector.reserve(num_underlying_devices);
+  for (int device_index = 0; device_index < num_underlying_devices;
+       ++device_index) {
+    underlying_devices_vector.push_back(underlying_devices[device_index]);
+  }
+  std::unique_ptr<ParallelDevice> parallel_device(
+      new ParallelDevice(underlying_devices_vector));
+  *device_info =
+      new NamedParallelDevice{device_name, std::move(parallel_device)};
+}
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device.h

@@ -0,0 +1,65 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_H_
+#define TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_H_
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+
+namespace tensorflow {
+namespace parallel_device {
+
+// Allocate a parallel device named `device_name` which forwards operations to
+// `underlying_devices`, maintaining "parallel tensors" with components placed
+// on each underlying device.
+//
+// For example if `device_name` is
+//   "/job:localhost/replica:0/task:0/device:CUSTOM:0"
+// and `underlying_devices` is
+//   {"/job:localhost/replica:0/task:0/device:GPU:0",
+//    "/job:localhost/replica:0/task:0/device:GPU:1"}
+// Then executing an operation on CUSTOM:0 will execute it on GPU:0 and GPU:1.
+//
+// Implicit copies onto `device_name` are allowed, replicating the value once
+// per device in `underlying_devices`. Implicit copies off of the device throw
+// an error.
+//
+// All component tensors must have the same dtype. Currently they must also have
+// the same shape, although this requirement may be relaxed in the future.
+//
+// `device_name` must not name an existing physical or custom device (see
+// the documentation for TFE_RegisterCustomDevice for more information).
+//
+// Tensors may be copied on or off the device explicitly using
+// TPUReplicatedInput and TPUReplicatedOutput respectively. For example, with
+// two component devices, running `x = TPUReplicatedInput(inputs=[a, b])` on the
+// parallel device creates a parallel tensor `x` with `a` on the first of
+// `underlying_devices` and `b` on the second. Running `a_unpacked, b_unpacked =
+// TPUReplicatedOutput(input=x, num_replicas=2)` un-packs the parallel tensor
+// into its components.
+//
+// The filled `device` struct and the allocated `device_info` struct may be
+// passed to TFE_RegisterCustomDevice. The `device_name` arguments must match.
+void AllocateParallelDevice(const char* device_name,
+                            const char* const* underlying_devices,
+                            int num_underlying_devices,
+                            TFE_CustomDevice* device, void** device_info);
+
+}  // namespace parallel_device
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_H_

tensorflow/c/eager/parallel_device/parallel_device_lib.cc

@@ -0,0 +1,649 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/eager/tfe_cancellation_manager_internal.h"
+#include "tensorflow/c/eager/tfe_op_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tsl/platform/thread_annotations.h"
+
+namespace tensorflow {
+namespace parallel_device {
+namespace {
+
+class OpDeleter {
+ public:
+  void operator()(TFE_Op* to_delete) const { TFE_DeleteOp(to_delete); }
+};
+
+using OpPtr = std::unique_ptr<TFE_Op, OpDeleter>;
+
+class StatusDeleter {
+ public:
+  void operator()(TF_Status* to_delete) const { TF_DeleteStatus(to_delete); }
+};
+
+using StatusPtr = std::unique_ptr<TF_Status, StatusDeleter>;
+
+class ExecutorDeleter {
+ public:
+  void operator()(TFE_Executor* to_delete) const {
+    TFE_DeleteExecutor(to_delete);
+  }
+};
+
+using ExecutorPtr = std::unique_ptr<TFE_Executor, ExecutorDeleter>;
+
+}  // namespace
+
+// Allows a single op at a time to be launched without blocking.
+//
+// DeviceThread itself is thread-safe, in that StartExecute will block if there
+// is a pending execution. Since StartExecute is equivalent to grabbing a lock,
+// multiple DeviceThreads should always be accessed in the same order to avoid
+// deadlocks.
+class DeviceThread {
+ public:
+  // Starts a background thread waiting for `StartExecute`.
+  explicit DeviceThread(const std::string& device, const bool is_async,
+                        const int in_flight_nodes_limit)
+      : status_(TF_NewStatus()),
+        // If the context's default exector is set to async, re-using that in
+        // each thread would cause collectives to deadlock. For consistency we
+        // create a new sync executor for every thread.
+        //
+        // TODO(allenl): We should have an async API that works with the
+        // parallel device.
+        device_(device),
+        executor_(
+            TFE_NewExecutor(is_async, /*enable_streaming_enqueue=*/true,
+                            /*in_flight_nodes_limit=*/in_flight_nodes_limit)),
+        op_(nullptr),
+        thread_(tensorflow::Env::Default()->StartThread(
+            tensorflow::ThreadOptions(), "parallel_device_execute",
+            std::bind(&DeviceThread::Run, this))) {}
+  ~DeviceThread();
+
+  // Requests that the worker thread execute the specified operation. Blocks
+  // until the previously pending operation (a StartExecute without a Join) has
+  // finished, if any.
+  //
+  // `cancellation_manager` must live until after `Join` finishes and pending
+  // `is_async` operations finish. In addition to allowing the caller to cancel
+  // the operation, its `StartCancel` method will be called if op execution
+  // fails on any device in order to cancel the others.
+  void StartExecute(TFE_Context* context, const char* operation_name,
+                    std::vector<TFE_TensorHandle*> inputs,
+                    const TFE_OpAttrs* attributes, int expected_max_outputs,
+                    CancellationManager& cancellation_manager,
+                    std::optional<int64_t> step_id = std::nullopt);
+  // Block until the previous `StartExecute` operation has executed. Forwards
+  // the status from `TFE_Execute` and returns outputs if the status is OK.
+  std::vector<TensorHandlePtr> Join(TF_Status* status);
+
+  // Block until all Ops finished running on the thread.
+  void AsyncWait(TF_Status* status);
+
+ private:
+  void Run();
+
+  void Execute(TFE_Context* context, const char* operation_name,
+               std::vector<TFE_TensorHandle*> inputs,
+               const TFE_OpAttrs* attributes, int expected_max_outputs,
+               std::vector<TensorHandlePtr>* outputs, TF_Status* status) const
+      TF_EXCLUSIVE_LOCKS_REQUIRED(execution_mutex_);
+
+  enum class ExecutionState {
+    kReadyToExecute,
+    kHasResult,
+    kIdle,
+    kShuttingDown,
+  };
+
+  tensorflow::mutex execution_mutex_;
+  ExecutionState execution_state_ TF_GUARDED_BY(execution_mutex_) =
+      ExecutionState::kIdle;
+  // Tells the worker thread that there is new work.
+  tensorflow::condition_variable start_execute_;
+  // The worker thread notifies that work has finished.
+  tensorflow::condition_variable finished_execute_;
+  // Notifies a StartExecute that the previous Join has finished.
+  tensorflow::condition_variable finished_join_;
+
+  // Temporary state between `StartExecute` and `Join`.
+  //
+  //   Inputs; pointers are to objects not owned by the DeviceThread, but which
+  //   are expected to live at least until `Join` finishes:
+  TFE_Context* context_ TF_GUARDED_BY(execution_mutex_);
+  const char* operation_name_ TF_GUARDED_BY(execution_mutex_);
+  std::optional<int64_t> step_id_ TF_GUARDED_BY(execution_mutex_) =
+      std::nullopt;
+  std::vector<TFE_TensorHandle*> op_inputs_ TF_GUARDED_BY(execution_mutex_);
+  const TFE_OpAttrs* attributes_ TF_GUARDED_BY(execution_mutex_);
+  int expected_max_outputs_ TF_GUARDED_BY(execution_mutex_);
+  CancellationManager* cancellation_manager_ TF_GUARDED_BY(execution_mutex_);
+  //   Outputs:
+  std::vector<TensorHandlePtr> op_outputs_ TF_GUARDED_BY(execution_mutex_);
+  // TF_Status is an incomplete type and so can't be stack allocated. To avoid
+  // unnecessary allocations each Execute call, we keep one heap-allocated
+  // version for the thread.
+  StatusPtr status_ TF_GUARDED_BY(execution_mutex_);
+
+  const std::string device_;
+  ExecutorPtr executor_ TF_GUARDED_BY(execution_mutex_);
+  mutable OpPtr op_ TF_GUARDED_BY(execution_mutex_);
+  std::unique_ptr<Thread> thread_;
+};
+
+DeviceThread::~DeviceThread() {
+  {
+    tensorflow::mutex_lock l(execution_mutex_);
+    execution_state_ = ExecutionState::kShuttingDown;
+  }
+  start_execute_.notify_one();
+}
+
+void DeviceThread::AsyncWait(TF_Status* status) {
+  tensorflow::mutex_lock l(execution_mutex_);
+  TFE_ExecutorWaitForAllPendingNodes(executor_.get(), status);
+  TFE_ExecutorClearError(executor_.get());
+}
+
+void DeviceThread::Run() {
+  while (true) {
+    {
+      tensorflow::mutex_lock l(execution_mutex_);
+      while (execution_state_ == ExecutionState::kIdle ||
+             execution_state_ == ExecutionState::kHasResult) {
+        start_execute_.wait(l);
+      }
+      if (execution_state_ == ExecutionState::kShuttingDown) {
+        return;
+      } else if (execution_state_ == ExecutionState::kReadyToExecute) {
+        // op_outputs_ may have been std::moved
+        op_outputs_ = std::vector<TensorHandlePtr>();
+        Execute(context_, operation_name_, std::move(op_inputs_), attributes_,
+                expected_max_outputs_, &op_outputs_, status_.get());
+        execution_state_ = ExecutionState::kHasResult;
+      }
+    }
+    finished_execute_.notify_one();
+  }
+}
+
+void DeviceThread::StartExecute(TFE_Context* context,
+                                const char* operation_name,
+                                std::vector<TFE_TensorHandle*> inputs,
+                                const TFE_OpAttrs* attributes,
+                                int expected_max_outputs,
+                                CancellationManager& cancellation_manager,
+                                std::optional<int64_t> step_id) {
+  {
+    tensorflow::mutex_lock l(execution_mutex_);
+    while (execution_state_ != ExecutionState::kIdle) {
+      // If there's already a pending execution, wait until Join finishes before
+      // starting on the next operation.
+      finished_join_.wait(l);
+    }
+    context_ = context;
+    operation_name_ = operation_name;
+    step_id_ = step_id;
+    op_inputs_ = inputs;
+    attributes_ = attributes;
+    expected_max_outputs_ = expected_max_outputs;
+    cancellation_manager_ = &cancellation_manager;
+    execution_state_ = ExecutionState::kReadyToExecute;
+  }
+  start_execute_.notify_one();
+}
+
+std::vector<TensorHandlePtr> DeviceThread::Join(TF_Status* status) {
+  std::vector<TensorHandlePtr> result;
+  {
+    tensorflow::mutex_lock l(execution_mutex_);
+    while (execution_state_ != ExecutionState::kHasResult) {
+      finished_execute_.wait(l);
+    }
+    if (TF_GetCode(status_.get()) != TF_OK) {
+      TF_SetStatus(status, TF_GetCode(status_.get()),
+                   TF_Message(status_.get()));
+      // Reset the member `status_` so future op executions (after recovery from
+      // the bad `status`) start with an OK status.
+      TF_SetStatus(status_.get(), TF_OK, "");
+    }
+    cancellation_manager_ = nullptr;
+    execution_state_ = ExecutionState::kIdle;
+    result = std::move(op_outputs_);
+  }
+  finished_join_.notify_one();
+  return result;
+}
+
+void DeviceThread::Execute(TFE_Context* context, const char* operation_name,
+                           std::vector<TFE_TensorHandle*> inputs,
+                           const TFE_OpAttrs* attributes,
+                           int expected_max_outputs,
+                           std::vector<TensorHandlePtr>* outputs,
+                           TF_Status* status) const {
+  if (op_ == nullptr) {
+    TFE_ContextSetExecutorForThread(context, executor_.get());
+    op_.reset(TFE_NewOp(context, operation_name, status));
+    if (TF_GetCode(status) != TF_OK) return;
+    TFE_OpSetDevice(op_.get(), device_.c_str(), status);
+    if (TF_GetCode(status) != TF_OK) return;
+  } else {
+    TFE_OpReset(op_.get(), operation_name, device_.c_str(), status);
+    if (TF_GetCode(status) != TF_OK) return;
+  }
+  TFE_OpAddAttrs(op_.get(), attributes);
+  for (int input_index = 0; input_index < inputs.size(); ++input_index) {
+    TFE_OpAddInput(op_.get(), inputs[input_index], status);
+    if (TF_GetCode(status) != TF_OK) return;
+  }
+  std::vector<TFE_TensorHandle*> unwrapped_results(expected_max_outputs);
+  int real_num_outputs = expected_max_outputs;
+  TFE_OpSetCancellationManager(op_.get(), wrap(cancellation_manager_), status);
+  if (TF_GetCode(status) != TF_OK) return;
+
+  // unwrap op_ and set step_id only if valid step id value was set.
+  // Currently only required for non-TFRT use cases, e.g., EagerOp.
+  if (step_id_.has_value()) {
+    tensorflow::unwrap(op_.get())->SetStepId(step_id_.value());
+  }
+
+  TFE_Execute(op_.get(), unwrapped_results.data(), &real_num_outputs, status);
+  if (TF_GetCode(status) != TF_OK) {
+    cancellation_manager_->StartCancel();
+    return;
+  }
+  unwrapped_results.resize(real_num_outputs);
+  outputs->reserve(real_num_outputs);
+  for (TFE_TensorHandle* unwrapped_result : unwrapped_results) {
+    outputs->emplace_back(unwrapped_result);
+  }
+}
+
+ParallelDevice::ParallelDevice(const std::vector<std::string>& devices,
+                               bool is_async, int in_flight_nodes_limit)
+    : underlying_devices_(devices),
+      default_cancellation_manager_(absl::make_unique<CancellationManager>()) {
+  device_threads_.reserve(devices.size());
+  for (int device_index = 0; device_index < devices.size(); ++device_index) {
+    device_threads_.emplace_back(new DeviceThread(
+        devices[device_index].c_str(), is_async, in_flight_nodes_limit));
+  }
+}
+
+// Necessary for a unique_ptr to a forward-declared type.
+ParallelDevice::~ParallelDevice() = default;
+
+std::unique_ptr<ParallelTensor> ParallelDevice::CopyToParallelDevice(
+    TFE_Context* context, TFE_TensorHandle* tensor, TF_Status* status) const {
+  std::vector<TensorHandlePtr> components;
+  components.reserve(underlying_devices_.size());
+  for (const std::string& underlying_device_name : underlying_devices_) {
+    TFE_TensorHandle* t = TFE_TensorHandleCopyToDevice(
+        tensor, context, underlying_device_name.c_str(), status);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+    components.emplace_back(t);
+  }
+  return ParallelTensor::FromTensorHandles(*this, std::move(components),
+                                           status);
+}
+
+std::unique_ptr<ParallelTensor> ParallelDevice::DeviceIDs(
+    TFE_Context* context, TF_Status* status) const {
+  std::vector<int32_t> ids;
+  ids.reserve(num_underlying_devices());
+  for (int i = 0; i < num_underlying_devices(); ++i) {
+    ids.push_back(i);
+  }
+  return ScalarsFromSequence<int32_t>(ids, context, status);
+}
+
+std::optional<std::vector<std::unique_ptr<ParallelTensor>>>
+ParallelDevice::Execute(TFE_Context* context,
+                        const std::vector<ParallelTensor*>& inputs,
+                        const char* operation_name,
+                        const TFE_OpAttrs* attributes, int expected_max_outputs,
+                        TF_Status* status) const {
+  std::vector<PartialTensorShape> expected_output_shapes(expected_max_outputs);
+  StartExecute(context, inputs, operation_name, attributes,
+               expected_max_outputs, *default_cancellation_manager_);
+  auto result = Join(expected_output_shapes, status);
+  if (TF_GetCode(status) != TF_OK) {
+    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> await_status(
+        TF_NewStatus(), TF_DeleteStatus);
+    // Wait until all pending nodes have completed since they may have a
+    // reference to default_cancellation_manager_. We ignore the status return
+    // since we already have a bad status to propagate.
+    TFE_ContextAsyncWait(context, await_status.get());
+    // Reset the cancellation manager on a bad status. Otherwise we'll cancel
+    // all future operations.
+    default_cancellation_manager_ = absl::make_unique<CancellationManager>();
+  }
+  return result;
+}
+
+void ParallelDevice::StartExecute(TFE_Context* context,
+                                  const std::vector<ParallelTensor*>& inputs,
+                                  const char* operation_name,
+                                  const TFE_OpAttrs* attributes,
+                                  int expected_max_outputs,
+                                  CancellationManager& cancellation_manager,
+                                  std::optional<int64_t> step_id) const {
+  for (int device_index = 0; device_index < underlying_devices_.size();
+       ++device_index) {
+    DeviceThread* device_thread = device_threads_[device_index].get();
+    std::vector<TFE_TensorHandle*> device_inputs;
+    device_inputs.reserve(inputs.size());
+    for (int input_index = 0; input_index < inputs.size(); ++input_index) {
+      // Parallel tensors are divided between operations by device.
+      device_inputs.push_back(inputs[input_index]->tensor(device_index));
+    }
+    device_thread->StartExecute(
+        context, operation_name, std::move(device_inputs), attributes,
+        expected_max_outputs, cancellation_manager, step_id);
+  }
+}
+
+void ParallelDevice::StartExecute(
+    TFE_Context* context,
+    const std::vector<std::vector<TFE_TensorHandle*>>& inputs,
+    const char* operation_name, const TFE_OpAttrs* attributes,
+    int expected_max_outputs, CancellationManager& cancellation_manager,
+    std::optional<int64_t> step_id) const {
+  for (int device_index = 0; device_index < underlying_devices_.size();
+       ++device_index) {
+    DeviceThread* device_thread = device_threads_[device_index].get();
+    std::vector<TFE_TensorHandle*> device_inputs;
+    device_inputs.reserve(inputs.size());
+    for (int input_index = 0; input_index < inputs.size(); ++input_index) {
+      // Parallel tensors are divided between operations by device.
+      device_inputs.push_back(inputs[input_index][device_index]);
+    }
+    device_thread->StartExecute(
+        context, operation_name, std::move(device_inputs), attributes,
+        expected_max_outputs, cancellation_manager, step_id);
+  }
+}
+
+void ParallelDevice::AsyncWait(TFE_Context* context, TF_Status* status) const {
+  StatusPtr first_bad_status(nullptr);
+
+  for (const auto& dt : device_threads_) {
+    StatusPtr async_wait_status(TF_NewStatus());
+    dt->AsyncWait(async_wait_status.get());
+    // Prefer non cancelled errors to uncover real failures.
+    if (TF_GetCode(async_wait_status.get()) != TF_OK &&
+        (first_bad_status == nullptr ||
+         TF_GetCode(first_bad_status.get()) == TF_CANCELLED)) {
+      first_bad_status.reset(TF_NewStatus());
+      TF_SetStatus(first_bad_status.get(), TF_GetCode(async_wait_status.get()),
+                   TF_Message(async_wait_status.get()));
+    }
+  }
+
+  if (first_bad_status != nullptr) {
+    TF_SetStatus(status, TF_GetCode(first_bad_status.get()),
+                 TF_Message(first_bad_status.get()));
+  }
+}
+
+std::optional<std::vector<std::unique_ptr<ParallelTensor>>>
+ParallelDevice::Join(
+    const std::vector<PartialTensorShape>& expected_output_shapes,
+    TF_Status* status) const {
+  std::optional<std::vector<std::unique_ptr<ParallelTensor>>> result;
+  // Compute per-device per-output tensors
+  std::vector<std::vector<TensorHandlePtr>> per_device_output_tensors;
+  per_device_output_tensors.reserve(underlying_devices_.size());
+  int first_op_output_count = 0;
+  StatusPtr first_bad_status(nullptr);
+  for (int device_index = 0; device_index < underlying_devices_.size();
+       ++device_index) {
+    DeviceThread* device_thread = device_threads_[device_index].get();
+    per_device_output_tensors.push_back(device_thread->Join(status));
+    // We will run every Join even if there are bad statuses in case the user
+    // wants to recover and continue running ops on the parallel device (which
+    // would otherwise deadlock).
+    if (TF_GetCode(status) != TF_OK &&
+        (first_bad_status == nullptr
+         // Prefer propagating non-cancellation related statuses to avoid
+         // shadowing the original failure.
+         || TF_GetCode(first_bad_status.get()) == TF_CANCELLED)) {
+      first_bad_status.reset(TF_NewStatus());
+      TF_SetStatus(first_bad_status.get(), TF_GetCode(status),
+                   TF_Message(status));
+    }
+
+    if (device_index == 0) {
+      first_op_output_count = per_device_output_tensors.rbegin()->size();
+    } else {
+      if (first_bad_status == nullptr &&
+          per_device_output_tensors.rbegin()->size() != first_op_output_count) {
+        first_bad_status.reset(TF_NewStatus());
+        TF_SetStatus(first_bad_status.get(), TF_INTERNAL,
+                     "Parallel ops produced different numbers of tensors.");
+      }
+    }
+  }
+  if (first_bad_status != nullptr) {
+    TF_SetStatus(status, TF_GetCode(first_bad_status.get()),
+                 TF_Message(first_bad_status.get()));
+    return result;
+  }
+  // For each output of the original operation, pack the per-device
+  // TensorHandles we've computed into a single parallel TensorHandle.
+  std::vector<std::unique_ptr<ParallelTensor>> per_device_outputs;
+  per_device_outputs.reserve(first_op_output_count);
+  for (int i = 0; i < first_op_output_count; ++i) {
+    std::vector<TensorHandlePtr> components;
+    components.reserve(underlying_devices_.size());
+    for (int j = 0; j < underlying_devices_.size(); ++j) {
+      components.push_back(std::move(per_device_output_tensors[j][i]));
+    }
+    if (expected_output_shapes[i].IsFullyDefined()) {
+      per_device_outputs.push_back(ParallelTensor::FromTensorHandles(
+          *this, std::move(components),
+          absl::Span<const int64_t>(expected_output_shapes[i].dim_sizes()),
+          status));
+    } else {
+      per_device_outputs.push_back(ParallelTensor::FromTensorHandles(
+          *this, std::move(components), status));
+    }
+    if (TF_GetCode(status) != TF_OK) return result;
+  }
+  result.emplace(std::move(per_device_outputs));
+  return result;
+}
+
+std::vector<std::string> ParallelDevice::SummarizeDeviceNames() const {
+  std::vector<DeviceNameUtils::ParsedName> parsed_components(
+      underlying_devices_.size());
+  for (int component_index = 0; component_index < underlying_devices_.size();
+       ++component_index) {
+    if (!DeviceNameUtils::ParseFullName(underlying_devices_[component_index],
+                                        &parsed_components[component_index]) ||
+        !DeviceNameUtils::IsSameAddressSpace(
+            underlying_devices_[component_index], underlying_devices_[0])) {
+      // Device names are from different address spaces, or we can't figure out
+      // whether they are, so we'll fully-qualify everything.
+      return underlying_devices_;
+    }
+  }
+  std::vector<std::string> local_names;
+  local_names.reserve(underlying_devices_.size());
+  for (const DeviceNameUtils::ParsedName& parsed_component :
+       parsed_components) {
+    local_names.push_back(
+        absl::StrCat(parsed_component.type, ":", parsed_component.id));
+  }
+  return local_names;
+}
+
+std::unique_ptr<ParallelTensor> ParallelTensor::FromTensorHandles(
+    const ParallelDevice& parallel_device,
+    std::vector<TensorHandlePtr> components, absl::Span<const int64_t> shape,
+    TF_Status* status) {
+  if (components.empty()) {
+    TF_SetStatus(status, TF_INTERNAL,
+                 "No components are provide for creating a ParallelTensor");
+    return nullptr;
+  }
+  TFE_TensorHandleGetStatus(components[0].get(), status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  TF_DataType dtype = TFE_TensorHandleDataType(components[0].get());
+  // Verify that the TensorHandle's shape and dtype match all of the component
+  // shapes and dtypes.
+  for (TensorHandlePtr& component : components) {
+    TFE_TensorHandleGetStatus(component.get(), status);
+    if (!status->status.ok()) {
+      return nullptr;
+    }
+    if (TFE_TensorHandleDataType(component.get()) != dtype) {
+      TF_SetStatus(status, TF_INTERNAL,
+                   "Components of a ParallelTensor must all have "
+                   "the same dtype");
+      return nullptr;
+    }
+  }
+  return std::unique_ptr<ParallelTensor>(
+      new ParallelTensor(parallel_device, std::move(components), shape, dtype));
+}
+
+std::unique_ptr<ParallelTensor> ParallelTensor::FromTensorHandles(
+    const ParallelDevice& parallel_device,
+    std::vector<TensorHandlePtr> components, TF_Status* status) {
+  if (components.empty()) {
+    TF_SetStatus(status, TF_INTERNAL,
+                 "No components are provided for creating a ParallelTensor");
+    return nullptr;
+  }
+  TFE_TensorHandleGetStatus(components[0].get(), status);
+  if (!status->status.ok()) {
+    return nullptr;
+  }
+
+  TF_DataType dtype = TFE_TensorHandleDataType(components[0].get());
+  // Verify that the combined TensorHandle's dtype matches all of the component
+  // dtypes.
+  for (TensorHandlePtr& component : components) {
+    TFE_TensorHandleGetStatus(component.get(), status);
+    if (!status->status.ok()) {
+      return nullptr;
+    }
+    if (TFE_TensorHandleDataType(component.get()) != dtype) {
+      TF_SetStatus(status, TF_INTERNAL,
+                   "Components of a ParallelTensor must all have "
+                   "the same dtype");
+      return nullptr;
+    }
+  }
+  return std::unique_ptr<ParallelTensor>(
+      new ParallelTensor(parallel_device, std::move(components), dtype));
+}
+
+absl::Status ParallelTensor::Shape(const std::vector<int64_t>** shape) const {
+  absl::MutexLock l(mu_);
+  if (!shape_.has_value()) {
+    TF_Status status;
+    PartialTensorShape combined_shape;
+    TF_RETURN_IF_ERROR(unwrap(tensors_[0].get())->Shape(&combined_shape));
+
+    for (const TensorHandlePtr& component : tensors_) {
+      PartialTensorShape component_shape;
+      TF_RETURN_IF_ERROR(unwrap(component.get())->Shape(&component_shape));
+      if (combined_shape.dims() < 0 ||
+          combined_shape.dims() != component_shape.dims()) {
+        PartialTensorShape first_shape;
+        TF_RETURN_IF_ERROR(unwrap(tensors_[0].get())->Shape(&first_shape));
+        return absl::UnimplementedError(absl::StrCat(
+            "Computing the shape of a ParallelTensor when the components do "
+            "not all have the same rank is not supported. One tensor had "
+            "shape ",
+            first_shape.DebugString(), " and another had shape ",
+            component_shape.DebugString()));
+      } else {
+        // Generalize differing axis lengths to "variable"/"unknown".
+        for (int axis_index = 0; axis_index < combined_shape.dims();
+             ++axis_index) {
+          int64_t axis_length = combined_shape.dim_size(axis_index);
+          if (axis_length != component_shape.dim_size(axis_index)) {
+            axis_length = -1;
+          }
+          TF_RETURN_IF_ERROR(
+              combined_shape.SetDimWithStatus(axis_index, axis_length));
+        }
+      }
+    }
+    auto dim_sizes = combined_shape.dim_sizes();
+    shape_ = std::vector<int64_t>(dim_sizes.begin(), dim_sizes.end());
+  }
+  *shape = &*shape_;
+  return absl::OkStatus();
+}
+
+absl::Status ParallelTensor::SummarizeValue(std::string& summary) {
+  summary = "{";
+  std::vector<std::string> summarized_devices = device_.SummarizeDeviceNames();
+  for (int component_index = 0; component_index < tensors_.size();
+       ++component_index) {
+    // TODO(allenl): Add a C API for summarizing tensors. Currently custom
+    // devices limiting themselves to a C API (for ABI compatibility) would need
+    // to implement summarization for component tensors themselves.
+    ImmediateExecutionTensorHandle* component =
+        tensorflow::unwrap(tensors_[component_index].get());
+    std::string component_summary;
+    TF_RETURN_IF_ERROR(component->SummarizeValue(component_summary));
+    absl::StrAppend(&summary, component_index == 0 ? "" : ", ", "\"",
+                    summarized_devices[component_index],
+                    "\": ", component_summary);
+  }
+  absl::StrAppend(&summary, "}");
+  return absl::OkStatus();
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device_lib.h

@@ -0,0 +1,304 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_LIB_H_
+#define TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_LIB_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "absl/types/variant.h"
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/tfe_op_internal.h"
+#include "tensorflow/c/safe_ptr.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace parallel_device {
+
+using TensorHandlePtr = tensorflow::Safe_TFE_TensorHandlePtr;
+
+class ParallelTensor;
+class DeviceThread;
+
+// Forwards operations to `devices`, maintaining ParallelTensor with components
+// placed on each underlying device.
+class ParallelDevice {
+ public:
+  // Eager async execution is only supported when remote eager is not in use
+  // (b/157523095).
+  explicit ParallelDevice(const std::vector<std::string>& devices,
+                          bool is_async = false, int in_flight_nodes_limit = 0);
+
+  ~ParallelDevice();
+
+  // Helper to copy a tensor handle from another device once for each component
+  // of the ParallelDevice.
+  //
+  // Sets a bad status and returns a nullptr if `tensor` is already on the
+  // ParallelDevice, or if the individual copies fail.
+  std::unique_ptr<ParallelTensor> CopyToParallelDevice(TFE_Context* context,
+                                                       TFE_TensorHandle* tensor,
+                                                       TF_Status* status) const;
+
+  // Construct a parallel tensor consisting of the scalar values from `values`.
+  template <typename DataType>
+  std::unique_ptr<ParallelTensor> ScalarsFromSequence(
+      absl::Span<const DataType> values, TFE_Context* context,
+      TF_Status* status) const;
+
+  // A parallel tensor with scalar integers numbering component devices.
+  std::unique_ptr<ParallelTensor> DeviceIDs(TFE_Context* context,
+                                            TF_Status* status) const;
+
+  // The number of devices operations run on.
+  size_t num_underlying_devices() const { return underlying_devices_.size(); }
+
+  // The devices operations run on.
+  const std::vector<std::string>& underlying_devices() const {
+    return underlying_devices_;
+  }
+
+  // Takes a description of a single operation being executed on the
+  // ParallelDevice, and in turn runs one operation per component device with
+  // its corresponding inputs from the input ParallelTensors. Wraps the
+  // resulting per-device and per-output TFE_TensorHandles into one
+  // ParallelTensor per output of the original operation.
+  //
+  // Attributes are forwarded to executed operations unmodified.
+  //
+  // The returned optional has a value if and only if `status` evaluates to
+  // TF_OK. Bad statuses are forwarded from underlying `TFE_Execute` calls, or
+  // if sanity checks on dtypes/metadata fail.
+  std::optional<std::vector<std::unique_ptr<ParallelTensor>>> Execute(
+      TFE_Context* context, const std::vector<ParallelTensor*>& inputs,
+      const char* operation_name, const TFE_OpAttrs* attributes,
+      int expected_max_outputs, TF_Status* status) const;
+
+  // A non-blocking version of `Execute`. After each call, `Join` must be called
+  // before `StartExecute` is called again. Using `StartExecute` with `Join`
+  // allows the caller to schedule computation on multiple ParallelDevices
+  // without sequencing those operations (first call `StartExecute` on each
+  // parallel device, then call `Join` on each; even if some of the `Join`s
+  // return a bad status the caller must run all of the `Join`s or any future
+  // `StartExecute`s will deadlock).
+  //
+  // If `is_async=false` (constructor argument), `cancellation_manager` must
+  // live until `Join` finishes. If `is_async=true` it must live until `Join` is
+  // followed by `TFE_ContextAsyncWait` to clear pending operations. It will be
+  // used to cancel all other operations if any fails.
+  //
+  // Set step_id to configure the step id used for rendezvous creation. step id
+  // of value -1 is reserved for global rendezvous and should not be set here.
+  void StartExecute(TFE_Context* context,
+                    const std::vector<ParallelTensor*>& inputs,
+                    const char* operation_name, const TFE_OpAttrs* attributes,
+                    int expected_max_outputs,
+                    CancellationManager& cancellation_manager,
+                    std::optional<int64_t> step_id = std::nullopt) const;
+
+  void StartExecute(TFE_Context* context,
+                    const std::vector<std::vector<TFE_TensorHandle*>>& inputs,
+                    const char* operation_name, const TFE_OpAttrs* attributes,
+                    int expected_max_outputs,
+                    CancellationManager& cancellation_manager,
+                    std::optional<int64_t> step_id = std::nullopt) const;
+
+  // Blocks until the previous `StartExecute` has run `TFE_Execute` on each
+  // device. If is_async=false (constructor argument) this means the ops have
+  // run and have results. If is_async=true it means that all of the
+  // device-specific executors have scheduled the op.
+  //
+  // Accepts inferred shapes for outputs (`expected_output_shapes`), which if
+  // fully defined will avoid querying the shapes of the underlying
+  // TensorHandles when ParallelTensor::Shape is called. This allows async
+  // computation to continue without blocking.
+  //
+  // The return status and value is the same as `Execute`.
+  std::optional<std::vector<std::unique_ptr<ParallelTensor>>> Join(
+      const std::vector<PartialTensorShape>& expected_output_shapes,
+      TF_Status* status) const;
+
+  void AsyncWait(TFE_Context* context, TF_Status* status) const;
+
+  // Device strings for component devices that only include a
+  // worker/task/replica if any of those differ across components. Useful for
+  // printing debug messages.
+  std::vector<std::string> SummarizeDeviceNames() const;
+
+ private:
+  // A sequence of device names, indicating which devices replicated operations
+  // are forwarded to.
+  const std::vector<std::string> underlying_devices_;
+  // A sequence of thread wrappers, one per device, for executing operations in
+  // parallel.
+  //
+  // Conceptually this is a thread pool with one thread per device. It requires
+  // less synchronization than a thread pool would for this task, since Execute
+  // acquires each thread in order (and so only one Execute will schedule
+  // blocking collective operations at a time), and avoids some dynamic
+  // allocation/scheduling.
+  //
+  // TODO(allenl): Keep a map from outer thread to list of inner threads rather
+  // than a single list of threads so aliased nested parallel devices don't
+  // re-use a thread.
+  std::vector<std::unique_ptr<DeviceThread>> device_threads_;
+  // A cancellation manager to use if the caller does not provide one. When ops
+  // are executed asynchronously this must outlive the queued op, so it can't be
+  // function-local to Execute.
+  mutable std::unique_ptr<CancellationManager> default_cancellation_manager_;
+};
+
+// Contains a tuple of tensors, one on each of the `underlying_devices_` of the
+// ParallelDevice.
+class ParallelTensor {
+ public:
+  // Construct a ParallelTensor from TensorHandles placed on the component
+  // devices of a ParallelDevice. If called, ParallelTensor::Shape inspects
+  // `components` to determine a shape.
+  static std::unique_ptr<ParallelTensor> FromTensorHandles(
+      const ParallelDevice& parallel_device,
+      std::vector<TensorHandlePtr> components, TF_Status* status);
+  // Uses the provided shape without additional checks, which avoids blocking
+  // when ParallelTensor::Shape is called.
+  static std::unique_ptr<ParallelTensor> FromTensorHandles(
+      const ParallelDevice& parallel_device,
+      std::vector<TensorHandlePtr> components, absl::Span<const int64_t> shape,
+      TF_Status* status);
+
+  size_t num_tensors() const { return tensors_.size(); }
+  TFE_TensorHandle* tensor(size_t index) const { return tensors_[index].get(); }
+
+  // If the `shape` argument to `FromTensorHandles` is specified, returns that.
+  //
+  // Otherwise if all of the tensors have the same shape, returns that via the
+  // `shape` output argument. This blocks waiting for async tensors, may return
+  // a delayed bad status encountered during async execution, and will return a
+  // bad status unless all tensors have the same shape.
+  absl::Status Shape(const std::vector<int64_t>** shape) const;
+  TF_DataType dtype() const { return dtype_; }
+
+  // Sets its output argument to a summary of the values of this tensor on every
+  // component device.
+  absl::Status SummarizeValue(std::string& summary);
+
+  std::vector<TensorHandlePtr> release_tensors() { return std::move(tensors_); }
+
+  std::vector<TFE_TensorHandle*> tensors() const {
+    std::vector<TFE_TensorHandle*> result;
+    result.reserve(tensors_.size());
+    for (const TensorHandlePtr& tensor : tensors_) {
+      result.emplace_back(tensor.get());
+    }
+    return result;
+  }
+
+ private:
+  ParallelTensor(const ParallelDevice& device,
+                 std::vector<TensorHandlePtr> tensors,
+                 absl::Span<const int64_t> shape, const TF_DataType dtype)
+      : device_(device),
+        tensors_(std::move(tensors)),
+        shape_(std::vector<int64_t>(shape.begin(), shape.end())),
+        dtype_(dtype) {}
+  ParallelTensor(const ParallelDevice& device,
+                 std::vector<TensorHandlePtr> tensors, const TF_DataType dtype)
+      : device_(device),
+        tensors_(std::move(tensors)),
+        shape_(std::nullopt),
+        dtype_(dtype) {}
+
+  const ParallelDevice& device_;
+  std::vector<TensorHandlePtr> tensors_;
+  // Parallel tensors are immutable but compute their shape lazily unless it is
+  // provided on construction. The optional has a value if the lazy computation
+  // has been completed or the shape was provided on construction.
+  mutable absl::Mutex mu_;
+  mutable std::optional<std::vector<int64_t>> shape_ ABSL_GUARDED_BY(mu_);
+  const TF_DataType dtype_;
+};
+
+template <typename DataType>
+std::unique_ptr<ParallelTensor> ParallelDevice::ScalarsFromSequence(
+    absl::Span<DataType const> values, TFE_Context* context,
+    TF_Status* status) const {
+  std::vector<TensorHandlePtr> components;
+  components.reserve(underlying_devices_.size());
+
+  if (values.size() != num_underlying_devices()) {
+    TF_SetStatus(
+        status, TF_INVALID_ARGUMENT,
+        "Number of values did not match number of underlying devices.");
+    return nullptr;
+  }
+  TF_DataType datatype_enum(
+      static_cast<TF_DataType>(DataTypeToEnum<DataType>().value));
+  for (int device_index = 0; device_index < num_underlying_devices();
+       ++device_index) {
+    auto device_value = absl::make_unique<DataType>();
+    *device_value = values[device_index];
+    std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tensor(
+        TF_NewTensor(
+            datatype_enum, /*dims=*/nullptr, /*num_dims=*/0,
+            device_value.release(), sizeof(DataType),
+            [](void* data, size_t, void* arg) {
+              delete reinterpret_cast<DataType*>(data);
+            },
+            nullptr),
+        TF_DeleteTensor);
+    // TODO(allenl): Here and when executing regular operations, we could hold
+    // on to one TFE_Op per device and just call TFE_ResetOp to avoid parsing
+    // device names repeatedly.
+    std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> const_op(
+        TFE_NewOp(context, "Const", status), TFE_DeleteOp);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+    TFE_OpSetDevice(const_op.get(), underlying_devices_[device_index].c_str(),
+                    status);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+    TFE_OpSetAttrTensor(const_op.get(), "value", tensor.get(), status);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+    TFE_OpSetAttrType(const_op.get(), "dtype", datatype_enum);
+    TFE_TensorHandle* device_handle;
+    int num_outputs = 1;
+    TFE_Execute(const_op.get(), &device_handle, &num_outputs, status);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+    components.emplace_back(device_handle);
+  }
+  return ParallelTensor::FromTensorHandles(*this, std::move(components),
+                                           status);
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_LIB_H_

tensorflow/c/eager/parallel_device/parallel_device_lib_test.cc

@@ -0,0 +1,340 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include "absl/status/status.h"
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_testlib.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/tf_buffer.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace parallel_device {
+
+using ::testing::ElementsAre;
+using ::testing::HasSubstr;
+
+TEST(PARALLEL_DEVICE_LIB, TestOpWithError) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*xla*/ false,
+          /* gpu_memory_allow_growth */ true, /* num_cpu_devices */
+          2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::vector<std::string> devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  ParallelDevice parallel_device(std::move(devices));
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> handle_op(
+      TFE_NewOp(context.get(), "VarHandleOp", status.get()), TFE_DeleteOp);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(handle_op.get(), "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(handle_op.get(), "shape", /*dims=*/nullptr, /*num_dims=*/0,
+                     status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  auto outputs =
+      parallel_device.Execute(context.get(), std::vector<ParallelTensor*>(),
+                              "VarHandleOp", TFE_OpGetAttrs(handle_op.get()),
+                              /*expected_max_outputs=*/1, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  const std::vector<std::unique_ptr<ParallelTensor>>& handles = *outputs;
+  std::vector<ParallelTensor*> handle_inputs;
+  handle_inputs.reserve(handles.size());
+  for (auto& handle : handles) {
+    handle_inputs.push_back(handle.get());
+  }
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> read_op(
+      TFE_NewOp(context.get(), "ReadVariableOp", status.get()), TFE_DeleteOp);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(read_op.get(), "dtype", TF_FLOAT);
+  parallel_device.Execute(context.get(), handle_inputs, "ReadVariableOp",
+                          TFE_OpGetAttrs(read_op.get()),
+                          /*expected_max_outputs=*/1, status.get());
+  ASSERT_FALSE(TF_GetCode(status.get()) == TF_OK);
+  TF_SetStatus(status.get(), TF_OK, "");
+
+  // Check that ops still run successfully on the device.
+  parallel_device.Execute(context.get(), std::vector<ParallelTensor*>(),
+                          "VarHandleOp", TFE_OpGetAttrs(handle_op.get()),
+                          /*expected_max_outputs=*/1, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+}
+
+TEST(PARALLEL_DEVICE_LIB, TestExplicitOutputShape) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*xla*/ false,
+          /* gpu_memory_allow_growth */ true, /* num_cpu_devices */
+          2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::vector<std::string> devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  ParallelDevice parallel_device(std::move(devices));
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> handle_op(
+      TFE_NewOp(context.get(), "VarHandleOp", status.get()), TFE_DeleteOp);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TFE_OpSetAttrType(handle_op.get(), "dtype", TF_FLOAT);
+  TFE_OpSetAttrShape(handle_op.get(), "shape", /*dims=*/nullptr, /*num_dims=*/0,
+                     status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  CancellationManager cancellation_manager;
+  parallel_device.StartExecute(context.get(), std::vector<ParallelTensor*>(),
+                               "VarHandleOp", TFE_OpGetAttrs(handle_op.get()),
+                               /*expected_max_outputs=*/1,
+                               cancellation_manager);
+  auto outputs = parallel_device.Join(
+      /*expected_output_shapes=*/{PartialTensorShape({})}, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  const std::vector<std::unique_ptr<ParallelTensor>>& handles = *outputs;
+  const std::vector<int64_t>* shape;
+  absl::Status s = handles[0]->Shape(&shape);
+  ASSERT_TRUE(s.ok());
+  EXPECT_EQ(0, shape->size());
+}
+
+TEST(PARALLEL_DEVICE_LIB, TestCancelOnError) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::vector<std::string> devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  ParallelDevice parallel_device(devices);
+  const FunctionDef assert_and_collective = FunctionDefHelper::Define(
+      // Name
+      "AssertAndCollective",
+      // Args
+      {"x: float", "condition: bool"},
+      // Return values
+      {"y: float"},
+      // Attr def
+      {},
+      // Nodes
+      {
+          {{"assert"},
+           "Assert",
+           {"condition", "x"},
+           {{"T", std::vector<DataType>{DT_FLOAT}}}},
+          {{"y"},
+           "CollectiveReduce",
+           {"x"},
+           {{"T", DT_FLOAT},
+            {"group_size", static_cast<int>(devices.size())},
+            {"group_key", 0},
+            {"instance_key", 0},
+            {"merge_op", "Add"},
+            {"final_op", "Id"},
+            {"subdiv_offsets", std::vector<int>()}},
+           /*dep=*/{"assert"}},
+      });
+  TF_ASSERT_OK(ContextFromInterface(unwrap(context.get()))
+                   ->AddFunctionDef(assert_and_collective));
+
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> call_op(
+      TFE_NewOp(context.get(), "AssertAndCollective", status.get()),
+      TFE_DeleteOp);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::unique_ptr<ParallelTensor> reduced_values =
+      parallel_device.ScalarsFromSequence<float>({1.0, 2.0}, context.get(),
+                                                 status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::unique_ptr<ParallelTensor> run_collective =
+      parallel_device.ScalarsFromSequence<bool>({true, true}, context.get(),
+                                                status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  auto outputs = parallel_device.Execute(
+      context.get(), {reduced_values.get(), run_collective.get()},
+      "AssertAndCollective", TFE_OpGetAttrs(call_op.get()),
+      /*expected_max_outputs=*/1, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(outputs->size(), 1);
+  ParallelTensor* parallel_result = (*outputs)[0].get();
+  ExpectScalarEq<float>(parallel_result->tensor(0), 3.);
+  ExpectScalarEq<float>(parallel_result->tensor(1), 3.);
+
+  run_collective = parallel_device.ScalarsFromSequence<bool>(
+      {true, false}, context.get(), status.get());
+  parallel_device.Execute(context.get(),
+                          {reduced_values.get(), run_collective.get()},
+                          "AssertAndCollective", TFE_OpGetAttrs(call_op.get()),
+                          /*expected_max_outputs=*/1, status.get());
+  EXPECT_NE(TF_GetCode(status.get()), TF_CANCELLED);
+  EXPECT_EQ(TF_GetCode(status.get()), TF_INVALID_ARGUMENT);
+  EXPECT_THAT(TF_Message(status.get()), HasSubstr("assertion failed"));
+
+  // Note that future collectives with the same context do not work at the
+  // moment; once canceled, the collective executor requires the program to be
+  // restarted / context to be reset.
+}
+
+TEST(PARALLEL_DEVICE_LIB, TestDifferentShapes) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*xla*/ false,
+          /* gpu_memory_allow_growth */ true, /* num_cpu_devices */
+          2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::vector<std::string> devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  ParallelDevice parallel_device(std::move(devices));
+  TensorHandlePtr two_vector = VectorFloatTensorHandle({3., 4.}, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TensorHandlePtr three_vector =
+      VectorFloatTensorHandle({5., 6., 7.}, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  std::vector<TensorHandlePtr> vector_handles;
+  vector_handles.reserve(2);
+  vector_handles.push_back(std::move(two_vector));
+  vector_handles.push_back(std::move(three_vector));
+  std::unique_ptr<ParallelTensor> unknown_length_vector =
+      ParallelTensor::FromTensorHandles(
+          parallel_device, std::move(vector_handles), status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  const std::vector<int64_t>* shape;
+  TF_ASSERT_OK(unknown_length_vector->Shape(&shape));
+  EXPECT_THAT(*shape, ElementsAre(-1));
+
+  TensorHandlePtr scalar = FloatTensorHandle(2., status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  two_vector = VectorFloatTensorHandle({3., 4.}, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  std::vector<TensorHandlePtr> mixed_handles;
+  mixed_handles.reserve(2);
+  mixed_handles.push_back(std::move(scalar));
+  mixed_handles.push_back(std::move(two_vector));
+  std::unique_ptr<ParallelTensor> unknown_dims_vector =
+      ParallelTensor::FromTensorHandles(parallel_device,
+                                        std::move(mixed_handles), status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK);
+  // Can't take the shape of a parallel tensor with varying numbers of axes, but
+  // running operations on them is OK.
+  TF_ASSERT_OK(unknown_length_vector->Shape(&shape));
+  EXPECT_THAT(*shape, ElementsAre(-1));
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> size_op(
+      TFE_NewOp(context.get(), "Size", status.get()), TFE_DeleteOp);
+  auto result = parallel_device.Execute(
+      context.get(), {unknown_dims_vector.get()}, "Size",
+      TFE_OpGetAttrs(size_op.get()), 1, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK);
+  TF_ASSERT_OK((*result)[0]->Shape(&shape));
+  EXPECT_EQ(0, shape->size());
+}
+
+TEST(PARALLEL_DEVICE_LIB, TestScalarsFromSequence) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  std::vector<std::string> devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  ParallelDevice parallel_device(std::move(devices));
+  {
+    std::unique_ptr<ParallelTensor> float_tensors =
+        parallel_device.ScalarsFromSequence<float>({10.0, 11.0}, context.get(),
+                                                   status.get());
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    ExpectScalarEq<float>(float_tensors->tensor(0), 10.0);
+    ExpectScalarEq<float>(float_tensors->tensor(1), 11.0);
+  }
+
+  {
+    std::unique_ptr<ParallelTensor> int_tensors =
+        parallel_device.ScalarsFromSequence<int>({5, 6}, context.get(),
+                                                 status.get());
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    ExpectScalarEq<int>(int_tensors->tensor(0), 5);
+    ExpectScalarEq<int>(int_tensors->tensor(1), 6);
+  }
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device_remote_test.cc

@@ -0,0 +1,158 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <array>
+#include <memory>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_testlib.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/distributed_runtime/master_env.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/cluster.pb.h"
+#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
+
+tensorflow::ServerDef GetServerDef(const std::string& job_name, int num_tasks) {
+  tensorflow::ServerDef server_def;
+  server_def.set_protocol("grpc");
+  server_def.set_job_name(job_name);
+  server_def.set_task_index(0);
+  tensorflow::ClusterDef* cluster_def = server_def.mutable_cluster();
+  tensorflow::JobDef* job_def = cluster_def->add_job();
+  job_def->set_name(job_name);
+  for (int i = 0; i < num_tasks; i++) {
+    int port = tensorflow::testing::PickUnusedPortOrDie();
+    job_def->mutable_tasks()->insert({i, absl::StrCat("localhost", ":", port)});
+    LOG(INFO) << "Picked test port: " << port << " for job: " << job_name
+              << ", task: " << i;
+  }
+  return server_def;
+}
+
+namespace tensorflow {
+namespace parallel_device {
+
+TEST(PARALLEL_DEVICE, TestRemoteBasic) {
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  tensorflow::ServerDef server_def = GetServerDef("worker", 3);
+
+  // This server def has the task index set to 0.
+  std::string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+
+  TFE_ContextSetServerDef(context.get(), 0, serialized.data(),
+                          serialized.size(), status.get());
+  EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+
+  BasicTestsForTwoDevices(context.get(),
+                          "/job:worker/replica:0/task:1/device:CPU:0",
+                          "/job:worker/replica:0/task:2/device:CPU:0");
+
+  worker_server1.release();
+  worker_server2.release();
+}
+
+TEST(PARALLEL_DEVICE, TestAsyncCopyOff) {
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  tensorflow::ServerDef server_def = GetServerDef("worker", 3);
+
+  // This server def has the task index set to 0.
+  std::string serialized = server_def.SerializeAsString();
+
+  server_def.set_task_index(1);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server1;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server1)
+                  .ok());
+  ASSERT_TRUE(worker_server1->Start().ok());
+
+  server_def.set_task_index(2);
+  std::unique_ptr<tensorflow::GrpcServer> worker_server2;
+  ASSERT_TRUE(tensorflow::GrpcServer::Create(
+                  server_def, tensorflow::Env::Default(), &worker_server2)
+                  .ok());
+  ASSERT_TRUE(worker_server2->Start().ok());
+
+  TFE_ContextSetServerDef(context.get(), 0, serialized.data(),
+                          serialized.size(), status.get());
+  EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
+
+  const char* first_device = "/job:worker/replica:0/task:1/device:CPU:0";
+  const char* second_device = "/job:worker/replica:0/task:2/device:CPU:0";
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{first_device, second_device};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr value_one(FloatTensorHandle(3., status.get()));
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  TensorHandlePtr value_two(FloatTensorHandle(-2., status.get()));
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  std::array<TFE_TensorHandle*, 2> in_components{value_one.get(),
+                                                 value_two.get()};
+  TensorHandlePtr combined_value = CreatePerDeviceValues(
+      context.get(), in_components, device_name, status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Loop to make synchronization failures more deterministic
+  for (int i = 0; i < 100; ++i) {
+    TensorHandlePtr multiply_result(
+        Multiply(context.get(), combined_value.get(), combined_value.get(),
+                 status.get()));
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+    std::array<TensorHandlePtr, 2> out_components;
+    ExtractPerDeviceValues(context.get(), multiply_result.get(),
+                           &out_components, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+    ExpectScalarEq<float>(out_components[0].get(), 9.);
+    ExpectScalarEq<float>(out_components[1].get(), 4.);
+  }
+
+  worker_server1.release();
+  worker_server2.release();
+}
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device_test.cc

@@ -0,0 +1,629 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <array>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_testlib.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/tf_buffer.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+// NOTE(allenl): These tests currently go through TFE_Execute and so are
+// integration testing rather than purely testing the parallel device. They
+// correspond fairly well to the implementation, but testing the C++ directly is
+// another option.
+
+namespace tensorflow {
+namespace parallel_device {
+
+using ::testing::HasSubstr;
+
+TEST(PARALLEL_DEVICE, TestBasicCPU) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  BasicTestsForTwoDevices(context.get(),
+                          "/job:localhost/replica:0/task:0/device:CPU:0",
+                          "/job:localhost/replica:0/task:0/device:CPU:1");
+}
+
+TEST(PARALLEL_DEVICE, TestBasicCPUAliased) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  BasicTestsForTwoDevices(context.get(),
+                          "/job:localhost/replica:0/task:0/device:CPU:0",
+                          "/job:localhost/replica:0/task:0/device:CPU:0");
+}
+
+TEST(PARALLEL_DEVICE, TestBasicTPUAliased) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Skip the test if no TPU is available.
+  std::unique_ptr<TF_DeviceList, decltype(&TF_DeleteDeviceList)> devices(
+      TFE_ContextListDevices(context.get(), status.get()), TF_DeleteDeviceList);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  bool has_tpu = false;
+  for (int device_index = 0; device_index < TF_DeviceListCount(devices.get());
+       ++device_index) {
+    std::string device_type =
+        TF_DeviceListType(devices.get(), device_index, status.get());
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    if (device_type == "TPU") {
+      has_tpu = true;
+      break;
+    }
+  }
+  if (has_tpu) {
+    BasicTestsForTwoDevices(context.get(),
+                            "/job:localhost/replica:0/task:0/device:TPU:0",
+                            "/job:localhost/replica:0/task:0/device:TPU:0");
+  }
+}
+
+TEST(PARALLEL_DEVICE, TestExplicitCopies) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  const char* first_device_name =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  const char* second_device_name =
+      "/job:localhost/replica:0/task:0/device:CPU:1";
+  std::array<const char*, 2> underlying_devices{first_device_name,
+                                                second_device_name};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr cpu_value(FloatTensorHandle(3., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Copying on to a parallel device must be explicit.
+  TensorHandlePtr failed_copy_on_result(TFE_TensorHandleCopyToDevice(
+      cpu_value.get(), context.get(), device_name, status.get()));
+  EXPECT_EQ(TF_GetCode(status.get()), TF_UNIMPLEMENTED);
+
+  std::array<TFE_TensorHandle*, 2> components{cpu_value.get(), cpu_value.get()};
+  TensorHandlePtr device_value = CreatePerDeviceValues(
+      context.get(), components, device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  // Copies off of parallel devices must be explicit.
+  TensorHandlePtr copy_off(TFE_TensorHandleCopyToDevice(
+      device_value.get(), context.get(), first_device_name, status.get()));
+  EXPECT_EQ(TF_GetCode(status.get()), TF_UNIMPLEMENTED);
+}
+
+TEST(PARALLEL_DEVICE, TestDifferentShapes) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Create two vectors with different lengths
+  std::vector<float> size_two_value{1., 2.};
+  std::vector<float> size_three_value{1., 2., 3.};
+  TensorHandlePtr size_two(
+      VectorFloatTensorHandle(size_two_value, status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  TensorHandlePtr size_three(
+      VectorFloatTensorHandle(size_three_value, status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Try to combine these values into a single parallel tensor.
+  std::array<TFE_TensorHandle*, 2> components{size_two.get(), size_three.get()};
+  TensorHandlePtr combined_value = CreatePerDeviceValues(
+      context.get(), components, device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  int num_axes = TFE_TensorHandleNumDims(combined_value.get(), status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  EXPECT_EQ(num_axes, 1);
+}
+
+TEST(PARALLEL_DEVICE, TestNestedParallelDevices) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/3),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Create a parallel device with two CPUs
+  const char* first_device_name =
+      "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> first_underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  RegisterParallelDevice(context.get(), first_device_name,
+                         first_underlying_devices, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Create a second parallel device with the first parallel device and one
+  // additional CPU.
+  const char* second_device_name =
+      "/job:localhost/replica:0/task:0/device:CUSTOM:1";
+  std::array<const char*, 2> second_underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CUSTOM:0",
+      "/job:localhost/replica:0/task:0/device:CPU:2"};
+  RegisterParallelDevice(context.get(), second_device_name,
+                         second_underlying_devices, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Create a tensor on the first parallel device
+  TensorHandlePtr value_one(FloatTensorHandle(1., status.get()));
+  TensorHandlePtr value_two(FloatTensorHandle(2., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::array<TFE_TensorHandle*, 2> components{value_one.get(), value_two.get()};
+  TensorHandlePtr first_combined_value = CreatePerDeviceValues(
+      context.get(), components, first_device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Nest the first parallel tensor into a second
+  TensorHandlePtr value_three(FloatTensorHandle(3., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  components[0] = first_combined_value.get();
+  components[1] = value_three.get();
+  TensorHandlePtr second_combined_value = CreatePerDeviceValues(
+      context.get(), components, second_device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr negative_one_cpu(FloatTensorHandle(3., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  components[0] = negative_one_cpu.get();
+  components[1] = negative_one_cpu.get();
+  TensorHandlePtr first_negative_one = CreatePerDeviceValues(
+      context.get(), components, first_device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  components[0] = first_negative_one.get();
+  components[1] = negative_one_cpu.get();
+  TensorHandlePtr second_negative_one = CreatePerDeviceValues(
+      context.get(), components, second_device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr multiply_result(
+      Multiply(context.get(), second_combined_value.get(),
+               second_negative_one.get(), status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Un-pack the parallel tensor to verify that the operation was
+  // successful. The resulting structure should be:
+  //   second_device{first_device{1. * 3., 2. * 3.}, 3. * 3.}.
+  std::array<TensorHandlePtr, 2> second_components;
+  ExtractPerDeviceValues(context.get(), multiply_result.get(),
+                         &second_components, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  ExpectScalarEq<float>(second_components[1].get(), 9.);
+
+  // Verify that the mirrors are placed on the component devices.
+  std::string first_device = TFE_TensorHandleBackingDeviceName(
+      second_components[0].get(), status.get());
+  ASSERT_EQ(second_underlying_devices[0], first_device);
+  std::string second_device = TFE_TensorHandleBackingDeviceName(
+      second_components[1].get(), status.get());
+  ASSERT_EQ(second_underlying_devices[1], second_device);
+
+  // Un-pack the first parallel device's tensor too
+  std::array<TensorHandlePtr, 2> first_components;
+  ExtractPerDeviceValues(context.get(), second_components[0].get(),
+                         &first_components, status.get());
+  ExpectScalarEq<float>(first_components[0].get(), 3.);
+  ExpectScalarEq<float>(first_components[1].get(), 6.);
+
+  first_device = TFE_TensorHandleBackingDeviceName(first_components[0].get(),
+                                                   status.get());
+  ASSERT_EQ(first_underlying_devices[0], first_device);
+  second_device = TFE_TensorHandleBackingDeviceName(first_components[1].get(),
+                                                    status.get());
+  ASSERT_EQ(first_underlying_devices[1], second_device);
+}
+
+TEST(PARALLEL_DEVICE, TestInvalidPacking) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 1> underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0"};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr value_one(FloatTensorHandle(1., status.get()));
+  TensorHandlePtr value_two(FloatTensorHandle(2., status.get()));
+  {
+    // Try to pack two TensorHandles onto a parallel device with a single
+    // component.
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+    std::array<TFE_TensorHandle*, 2> components{value_one.get(),
+                                                value_two.get()};
+    TensorHandlePtr combined_value = CreatePerDeviceValues(
+        context.get(), components, device_name, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+        << TF_Message(status.get());
+  }
+
+  {
+    // Try to extract the wrong number of components from a parallel tensor
+    std::array<TFE_TensorHandle*, 1> correct_components{value_one.get()};
+    TensorHandlePtr combined_value = CreatePerDeviceValues(
+        context.get(), correct_components, device_name, status.get());
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+    std::array<TensorHandlePtr, 2> incorrect_components;
+    ExtractPerDeviceValues(context.get(), combined_value.get(),
+                           &incorrect_components, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+        << TF_Message(status.get());
+  }
+
+  {
+    // Try to pass a ParallelTensor to TPUReplicatedInput
+    std::array<TFE_TensorHandle*, 1> correct_components{value_one.get()};
+    TensorHandlePtr combined_value = CreatePerDeviceValues(
+        context.get(), correct_components, device_name, status.get());
+    ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+    std::array<TFE_TensorHandle*, 1> incorrect_components{combined_value.get()};
+    TensorHandlePtr recombined_value = CreatePerDeviceValues(
+        context.get(), incorrect_components, device_name, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+        << TF_Message(status.get());
+  }
+
+  {
+    // Try to pass a non-parallel tensor to TPUReplicatedOutput
+    std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+        TFE_NewOp(context.get(), "TPUReplicatedOutput", status.get()),
+        TFE_DeleteOp);
+    if (TF_GetCode(status.get()) != TF_OK) return;
+    TFE_OpSetAttrInt(op.get(), "num_replicas", 1);
+    TFE_OpAddInput(op.get(), value_one.get(), status.get());
+    if (TF_GetCode(status.get()) != TF_OK) return;
+    TFE_OpSetDevice(op.get(), device_name, status.get());
+    if (TF_GetCode(status.get()) != TF_OK) return;
+
+    TFE_TensorHandle* result_handles;
+    int num_retvals = 1;
+    TFE_Execute(op.get(), &result_handles, &num_retvals, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+        << TF_Message(status.get());
+  }
+
+  {
+    // Try to pass 0 inputs to TPUReplicatedOutput
+    std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+        TFE_NewOp(context.get(), "TPUReplicatedOutput", status.get()),
+        TFE_DeleteOp);
+    if (TF_GetCode(status.get()) != TF_OK) return;
+    TFE_OpSetAttrInt(op.get(), "num_replicas", 1);
+    TFE_OpSetDevice(op.get(), device_name, status.get());
+    if (TF_GetCode(status.get()) != TF_OK) return;
+
+    TFE_TensorHandle* result_handles;
+    int num_retvals = 1;
+    TFE_Execute(op.get(), &result_handles, &num_retvals, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_INVALID_ARGUMENT)
+        << TF_Message(status.get());
+  }
+}
+
+TensorHandlePtr CollectiveSum(TFE_Context* context, TFE_TensorHandle* input,
+                              int group_size, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "CollectiveReduce", status), TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  const char* device = TFE_TensorHandleDeviceName(input, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op.get(), "T", TFE_TensorHandleDataType(input));
+  TFE_OpSetAttrInt(op.get(), "group_size", group_size);
+  TFE_OpSetAttrInt(op.get(), "group_key", 0);
+  TFE_OpSetAttrInt(op.get(), "instance_key", 0);
+  const std::string merge_op("Add");
+  TFE_OpSetAttrString(op.get(), "merge_op", merge_op.c_str(),
+                      merge_op.length());
+  const std::string final_op("Id");
+  TFE_OpSetAttrString(op.get(), "final_op", final_op.c_str(),
+                      final_op.length());
+  TFE_OpSetAttrIntList(op.get(), "subdiv_offsets", nullptr, 0);
+
+  TFE_OpAddInput(op.get(), input, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  TFE_TensorHandle* result_handle;
+  int num_retvals = 1;
+  TFE_Execute(op.get(), &result_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  return TensorHandlePtr(result_handle);
+}
+
+void TestCollective(bool async) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  TFE_ContextOptionsSetAsync(opts.get(), async);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Create a tensor on the parallel device
+  TensorHandlePtr value_one(FloatTensorHandle(1., status.get()));
+  TensorHandlePtr value_two(FloatTensorHandle(2., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::array<TFE_TensorHandle*, 2> components{value_one.get(), value_two.get()};
+  TensorHandlePtr parallel_value = CreatePerDeviceValues(
+      context.get(), components, device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  // Run a collective sum, so each component should now be the same.
+  TensorHandlePtr reduced(
+      CollectiveSum(context.get(), parallel_value.get(), 2, status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  std::array<TensorHandlePtr, 2> result_components;
+  ExtractPerDeviceValues(context.get(), reduced.get(), &result_components,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  ExpectScalarEq<float>(result_components[0].get(), 3.);
+  ExpectScalarEq<float>(result_components[1].get(), 3.);
+}
+
+TEST(PARALLEL_DEVICE, TestCollectiveSync) { TestCollective(/*async=*/false); }
+
+// Note that ops on the parallel device currently don't execute
+// asynchronously. The test is just that we don't get deadlocks.
+TEST(PARALLEL_DEVICE, TestCollectiveAsync) { TestCollective(/*async=*/true); }
+
+void RegisterCollectiveMulFunction(TFE_Context* context,
+                                   const char* function_name, int group_size,
+                                   TF_Status* status) {
+  std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)> body(TF_NewGraph(),
+                                                            TF_DeleteGraph);
+  TF_OperationDescription* placeholder_desc =
+      TF_NewOperation(body.get(), "Placeholder", "Placeholder");
+  TF_SetAttrType(placeholder_desc, "dtype", TF_FLOAT);
+  TF_Operation* placeholder_op = TF_FinishOperation(placeholder_desc, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TF_Output x{placeholder_op, 0};
+
+  TF_OperationDescription* reduce_desc =
+      TF_NewOperation(body.get(), "CollectiveReduce", "CollectiveReduce");
+  TF_SetAttrType(reduce_desc, "T", TF_FLOAT);
+  TF_SetAttrInt(reduce_desc, "group_size", group_size);
+  TF_SetAttrInt(reduce_desc, "group_key", 0);
+  TF_SetAttrInt(reduce_desc, "instance_key", 0);
+
+  const std::string merge_op("Mul");
+  TF_SetAttrString(reduce_desc, "merge_op", merge_op.c_str(),
+                   merge_op.length());
+  const std::string final_op("Id");
+  TF_SetAttrString(reduce_desc, "final_op", final_op.c_str(),
+                   final_op.length());
+  TF_SetAttrIntList(reduce_desc, "subdiv_offsets", nullptr, 0);
+  TF_AddInput(reduce_desc, x);
+  TF_Operation* reduce_op = TF_FinishOperation(reduce_desc, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TF_Operation* operations[]{placeholder_op, reduce_op};
+  TF_Output y{reduce_op, 0};
+  const char* output_name = "y";
+  std::unique_ptr<TF_Function, decltype(&TF_DeleteFunction)> function(
+      TF_GraphToFunction(
+          /* fn_body */ body.get(), /* fn_name */ function_name,
+          /* append_hash_to_fn_name */ 0, /* num_opers */ 2,
+          /* opers */ operations, /* ninputs */ 1, /* inputs */ &x,
+          /* noutputs */ 1, /* outputs */ &y, /* output_names */ &output_name,
+          /* opts */ nullptr, /* description */ "", /* status */ status),
+      TF_DeleteFunction);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_ContextAddFunction(context, function.get(), status);
+}
+
+TEST(PARALLEL_DEVICE, TestFunction) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* function_name = "test_reduce_mul";
+  RegisterCollectiveMulFunction(context.get(), function_name, 2, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TensorHandlePtr value_one(FloatTensorHandle(7., status.get()));
+  TensorHandlePtr value_two(FloatTensorHandle(9., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::array<TFE_TensorHandle*, 2> components{value_one.get(), value_two.get()};
+  TensorHandlePtr parallel_value = CreatePerDeviceValues(
+      context.get(), components, device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context.get(), function_name, status.get()), TFE_DeleteOp);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  TFE_OpSetDevice(op.get(), device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  TFE_OpAddInput(op.get(), parallel_value.get(), status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  TFE_TensorHandle* raw_result_handle;
+  int num_retvals = 1;
+  TFE_Execute(op.get(), &raw_result_handle, &num_retvals, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  TensorHandlePtr reduced(raw_result_handle);
+
+  std::array<TensorHandlePtr, 2> result_components;
+  ExtractPerDeviceValues(context.get(), reduced.get(), &result_components,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  ExpectScalarEq<float>(result_components[0].get(), 7. * 9.);
+  ExpectScalarEq<float>(result_components[1].get(), 7. * 9.);
+
+  std::string first_device = TFE_TensorHandleBackingDeviceName(
+      result_components[0].get(), status.get());
+  ASSERT_EQ(underlying_devices[0], first_device);
+  std::string second_device = TFE_TensorHandleBackingDeviceName(
+      result_components[1].get(), status.get());
+  ASSERT_EQ(underlying_devices[1], second_device);
+}
+
+TEST(PARALLEL_DEVICE, TestSummaryString) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TFE_ContextOptions, decltype(&TFE_DeleteContextOptions)> opts(
+      TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> config(
+      TF_CreateConfig(
+          /*enable_xla_compilation=*/false,
+          /*gpu_memory_allow_growth=*/true, /*num_cpu_devices=*/2),
+      TF_DeleteBuffer);
+  TFE_ContextOptionsSetConfig(opts.get(), config->data, config->length,
+                              status.get());
+  std::unique_ptr<TFE_Context, decltype(&TFE_DeleteContext)> context(
+      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{
+      "/job:localhost/replica:0/task:0/device:CPU:0",
+      "/job:localhost/replica:0/task:0/device:CPU:1"};
+  RegisterParallelDevice(context.get(), device_name, underlying_devices,
+                         status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  TensorHandlePtr cpu_value(FloatTensorHandle(3., status.get()));
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  std::array<TFE_TensorHandle*, 2> components{cpu_value.get(), cpu_value.get()};
+  TensorHandlePtr device_value = CreatePerDeviceValues(
+      context.get(), components, device_name, status.get());
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  ImmediateExecutionTensorHandle* unwrapped_handle =
+      tensorflow::unwrap(device_value.get());
+  std::string summarized;
+  TF_ASSERT_OK(unwrapped_handle->SummarizeValue(summarized));
+  EXPECT_THAT(summarized, HasSubstr("\"CPU:0\": 3"));
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device_testlib.cc

@@ -0,0 +1,300 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/eager/parallel_device/parallel_device_testlib.h"
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/test.h"
+
+// NOTE(allenl): These tests currently go through TFE_Execute and so are
+// integration testing rather than purely testing the parallel device. They
+// correspond fairly well to the implementation, but testing the C++ directly is
+// another option.
+
+namespace tensorflow {
+namespace parallel_device {
+
+Variable* Variable::Create(TFE_Context* context, TF_DataType type,
+                           const int64_t* dims, const int num_dims,
+                           const char* device, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "VarHandleOp", status), TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op.get(), "dtype", type);
+  TFE_OpSetAttrShape(op.get(), "shape", dims, num_dims, status);
+  TFE_OpSetAttrString(op.get(), "container", "", 0);
+  // Use the special GUID for no buffer sharing
+  //
+  // TODO(allenl): Should we provide a better API for this? AFAIK this is the
+  // only reasonable way to make variables with no aliasing using the eager C
+  // API.
+  std::string no_sharing = "cd2c89b7-88b7-44c8-ad83-06c2a9158347";
+  TFE_OpSetAttrString(op.get(), "shared_name", no_sharing.c_str(),
+                      no_sharing.length());
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_TensorHandle* var_handle = nullptr;
+  int num_retvals = 1;
+  TFE_Execute(op.get(), &var_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  return new Variable(var_handle, type);
+}
+
+void Variable::Destroy(TFE_Context* context, TF_Status* status) {
+  // Free the backing buffer for the variable.
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "DestroyResourceOp", status), &TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpAddInput(op.get(), handle_, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  const char* device = TFE_TensorHandleDeviceName(handle_, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  int num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  // Delete the variable handle itself.
+  TFE_DeleteTensorHandle(handle_);
+}
+
+TensorHandlePtr Variable::Read(TFE_Context* context, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "ReadVariableOp", status), &TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpAddInput(op.get(), handle_, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  const char* device = TFE_TensorHandleDeviceName(handle_, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrType(op.get(), "dtype", type_);
+  int num_retvals = 1;
+  TFE_TensorHandle* var_value = nullptr;
+  TFE_Execute(op.get(), &var_value, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  return TensorHandlePtr(var_value);
+}
+
+void Variable::GeneralAssignment(const char* op_name, TFE_Context* context,
+                                 TFE_TensorHandle* value, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, op_name, status), &TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpSetAttrType(op.get(), "dtype", type_);
+  TFE_OpAddInput(op.get(), handle_, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpAddInput(op.get(), value, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  const char* device = TFE_TensorHandleDeviceName(handle_, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpSetDevice(op.get(), device, status);
+
+  int num_retvals = 0;
+  TFE_Execute(op.get(), nullptr, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return;
+}
+
+void Variable::AssignAdd(TFE_Context* context, TFE_TensorHandle* value,
+                         TF_Status* status) {
+  GeneralAssignment("AssignAddVariableOp", context, value, status);
+}
+
+void Variable::Assign(TFE_Context* context, TFE_TensorHandle* value,
+                      TF_Status* status) {
+  GeneralAssignment("AssignVariableOp", context, value, status);
+}
+
+// Passed to `TF_NewTensor` to indicate how an array of floats should be
+// deleted.
+static void FloatDeallocator(void* data, size_t, void* arg) {
+  delete[] static_cast<float*>(data);
+}
+
+// Creates a TFE_TensorHandle with value `v`.
+TensorHandlePtr FloatTensorHandle(float v, TF_Status* status) {
+  const int num_bytes = sizeof(float);
+  float* values = new float[1];
+  values[0] = v;
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tensor(
+      TF_NewTensor(TF_FLOAT, nullptr, 0, values, num_bytes, &FloatDeallocator,
+                   nullptr),
+      TF_DeleteTensor);
+  return TensorHandlePtr(TFE_NewTensorHandle(tensor.get(), status));
+}
+
+// Creates a rank-one TFE_TensorHandle with value `v`.
+TensorHandlePtr VectorFloatTensorHandle(const std::vector<float>& v,
+                                        TF_Status* status) {
+  const int num_bytes = v.size() * sizeof(float);
+  float* values = new float[v.size()];
+  memcpy(values, v.data(), num_bytes);
+  int64_t dims = v.size();
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tensor(
+      TF_NewTensor(TF_FLOAT, &dims, 1 /* num_dims */, values, num_bytes,
+                   &FloatDeallocator, nullptr),
+      TF_DeleteTensor);
+  return TensorHandlePtr(TFE_NewTensorHandle(tensor.get(), status));
+}
+
+// Helper to un-pack `num_replicas` TFE_TensorHandles from one parallel handle.
+template <std::size_t num_replicas>
+void ExtractPerDeviceValues(
+    TFE_Context* context, TFE_TensorHandle* input,
+    std::array<TensorHandlePtr, num_replicas>* components, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "TPUReplicatedOutput", status), TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpSetAttrInt(op.get(), "num_replicas", num_replicas);
+  TFE_OpAddInput(op.get(), input, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  const char* device = TFE_TensorHandleDeviceName(input, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return;
+
+  TFE_TensorHandle* result_handles[num_replicas];
+  int num_retvals = num_replicas;
+  TFE_Execute(op.get(), result_handles, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return;
+  for (int i = 0; i < num_replicas; ++i) {
+    (*components)[i].reset(result_handles[i]);
+  }
+}
+
+TensorHandlePtr Multiply(TFE_Context* context, TFE_TensorHandle* first,
+                         TFE_TensorHandle* second, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "Mul", status), TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpAddInput(op.get(), first, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpAddInput(op.get(), second, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  const char* first_device = TFE_TensorHandleDeviceName(first, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetDevice(op.get(), first_device, status);
+
+  TFE_TensorHandle* result_handle;
+  int num_retvals = 1;
+  TFE_Execute(op.get(), &result_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  return TensorHandlePtr(result_handle);
+}
+
+// Create and modify a variable placed on a parallel device which composes
+// `first_device` and `second_device`.
+void BasicTestsForTwoDevices(TFE_Context* context, const char* first_device,
+                             const char* second_device) {
+  // Register the custom device
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
+  std::array<const char*, 2> underlying_devices{first_device, second_device};
+  RegisterParallelDevice(context, device_name, underlying_devices,
+                         status.get());
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Create a variable handle (uninitialized to start) placed on the parallel
+  // device.
+  std::function<void(Variable*)> variable_deleter = [&](Variable* to_delete) {
+    to_delete->Destroy(context, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+    delete to_delete;
+  };
+  std::unique_ptr<Variable, decltype(variable_deleter)> variable(
+      Variable::Create(context, TF_FLOAT, /* Scalar */ {}, 0, device_name,
+                       status.get()),
+      variable_deleter);
+  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+  // Assign an initial value to the variable, mirroring it to each component
+  // device.
+  {
+    TensorHandlePtr initial_value_cpu = FloatTensorHandle(20., status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+    std::array<TFE_TensorHandle*, 2> components{initial_value_cpu.get(),
+                                                initial_value_cpu.get()};
+    TensorHandlePtr initial_value =
+        CreatePerDeviceValues(context, components, device_name, status.get());
+    variable->Assign(context, initial_value.get(), status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+  }
+
+  // Read from the variable and verify that we have a parallel tensor.
+  {
+    TensorHandlePtr read = variable->Read(context, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+    std::array<TensorHandlePtr, 2> components;
+    ExtractPerDeviceValues(context, read.get(), &components, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+    ExpectScalarEq<float>(components[0].get(), 20.);
+    ExpectScalarEq<float>(components[1].get(), 20.);
+
+    std::string first_device =
+        TFE_TensorHandleBackingDeviceName(components[0].get(), status.get());
+    ASSERT_EQ(underlying_devices[0], first_device);
+    std::string second_device =
+        TFE_TensorHandleBackingDeviceName(components[1].get(), status.get());
+    ASSERT_EQ(underlying_devices[1], second_device);
+  }
+
+  // Add a parallel tensor with different values on each device to the variable.
+  {
+    TensorHandlePtr value_one(FloatTensorHandle(3., status.get()));
+    TensorHandlePtr value_two(FloatTensorHandle(-2., status.get()));
+    std::array<TFE_TensorHandle*, 2> components{value_one.get(),
+                                                value_two.get()};
+    TensorHandlePtr combined_value =
+        CreatePerDeviceValues(context, components, device_name, status.get());
+    variable->AssignAdd(context, combined_value.get(), status.get());
+  }
+
+  // Read the variable and verify that each component has the right modified
+  // value.
+  {
+    TensorHandlePtr read = variable->Read(context, status.get());
+    std::array<TensorHandlePtr, 2> components;
+    ExtractPerDeviceValues(context, read.get(), &components, status.get());
+    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
+
+    ExpectScalarEq<float>(components[0].get(), 23.);
+    ExpectScalarEq<float>(components[1].get(), 18.);
+
+    std::string first_device =
+        TFE_TensorHandleBackingDeviceName(components[0].get(), status.get());
+    ASSERT_EQ(underlying_devices[0], first_device);
+    std::string second_device =
+        TFE_TensorHandleBackingDeviceName(components[1].get(), status.get());
+    ASSERT_EQ(underlying_devices[1], second_device);
+  }
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow

tensorflow/c/eager/parallel_device/parallel_device_testlib.h

@@ -0,0 +1,172 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_TESTLIB_H_
+#define TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_TESTLIB_H_
+
+#include <array>
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device_lib.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace parallel_device {
+
+// A helper for performing common operations on variables. A much more
+// restricted stand-in for tf.Variable in Python.
+class Variable {
+ public:
+  // Construct a Variable from a resource-dtype TFE_TensorHandle and an
+  // indication of the dtype of the variable's value.
+  //
+  // Note that creating this resource-dtype handle can fail, so `Create` is a
+  // separate static method which returns a status.
+  Variable(TFE_TensorHandle* handle, TF_DataType type)
+      : handle_(handle), type_(type) {}
+
+  // Helper for constructing a resource handle and wrapping it in a `Variable`
+  // object.
+  static Variable* Create(TFE_Context* context, TF_DataType type,
+                          const int64_t* dims, const int num_dims,
+                          const char* device, TF_Status* status);
+  // Dereferences the backing buffer for the variable. Note that since this can
+  // fail (it runs operations), it must be called explicitly and the resulting
+  // `status` checked.
+  void Destroy(TFE_Context* context, TF_Status* status);
+
+  // Reads from the variable.
+  TensorHandlePtr Read(TFE_Context* context, TF_Status* status);
+  // Assigns a new value to the variable.
+  void Assign(TFE_Context* context, TFE_TensorHandle* value, TF_Status* status);
+  // Adds `value` to the existing value of the variable.
+  void AssignAdd(TFE_Context* context, TFE_TensorHandle* value,
+                 TF_Status* status);
+
+ private:
+  // Helper for running any single-argument assignment ops (Assign, AssignAdd,
+  // AssignSub, ...).
+  void GeneralAssignment(const char* op_name, TFE_Context* context,
+                         TFE_TensorHandle* value, TF_Status* status);
+
+  // The a handle for the resource-dtype tensor pointing to the variable's
+  // buffer.
+  TFE_TensorHandle* handle_;
+  // The dtype of the variable's buffer (input dtype for assignments, output
+  // dtype of read operations).
+  TF_DataType type_;
+};
+
+// Creates a TFE_TensorHandle with value `v`.
+TensorHandlePtr FloatTensorHandle(float v, TF_Status* status);
+
+// Creates a rank-one TFE_TensorHandle with value `v`.
+TensorHandlePtr VectorFloatTensorHandle(const std::vector<float>& v,
+                                        TF_Status* status);
+
+// Helper to un-pack `num_replicas` TFE_TensorHandles from one parallel handle.
+template <std::size_t num_replicas>
+void ExtractPerDeviceValues(
+    TFE_Context* context, TFE_TensorHandle* input,
+    std::array<TensorHandlePtr, num_replicas>* components, TF_Status* status);
+
+// Helper to pack `num_replicas` TFE_TensorHandles into one parallel handle.
+template <std::size_t num_replicas>
+TensorHandlePtr CreatePerDeviceValues(
+    TFE_Context* context,
+    const std::array<TFE_TensorHandle*, num_replicas>& components,
+    const char* device, TF_Status* status);
+
+TensorHandlePtr Multiply(TFE_Context* context, TFE_TensorHandle* first,
+                         TFE_TensorHandle* second, TF_Status* status);
+
+// Assert that `handle` is equal to `expected_value`.
+template <typename value_type>
+void ExpectScalarEq(TFE_TensorHandle* handle, value_type expected_value);
+
+template <std::size_t num_devices>
+void RegisterParallelDevice(
+    TFE_Context* context, const char* device_name,
+    const std::array<const char*, num_devices>& underlying_devices,
+    TF_Status* status);
+
+// Create and modify a variable placed on a parallel device which composes
+// `first_device` and `second_device`.
+void BasicTestsForTwoDevices(TFE_Context* context, const char* first_device,
+                             const char* second_device);
+
+// Implementations of templated functions ******************************
+
+template <std::size_t num_replicas>
+TensorHandlePtr CreatePerDeviceValues(
+    TFE_Context* context,
+    const std::array<TFE_TensorHandle*, num_replicas>& components,
+    const char* device, TF_Status* status) {
+  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
+      TFE_NewOp(context, "TPUReplicatedInput", status), TFE_DeleteOp);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  TFE_OpSetAttrInt(op.get(), "N", num_replicas);
+  for (int i = 0; i < num_replicas; ++i) {
+    TFE_OpAddInput(op.get(), components[i], status);
+    if (TF_GetCode(status) != TF_OK) return nullptr;
+  }
+  TFE_OpSetDevice(op.get(), device, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+
+  TFE_TensorHandle* result_handle;
+  int num_retvals = 1;
+  TFE_Execute(op.get(), &result_handle, &num_retvals, status);
+  if (TF_GetCode(status) != TF_OK) return nullptr;
+  return TensorHandlePtr(result_handle);
+}
+
+template <typename value_type>
+void ExpectScalarEq(TFE_TensorHandle* handle, value_type expected_value) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> actual_value(
+      TFE_TensorHandleResolve(handle, status.get()), TF_DeleteTensor);
+  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
+  ASSERT_EQ(TF_TensorType(actual_value.get()),
+            static_cast<TF_DataType>(DataTypeToEnum<value_type>().value));
+  EXPECT_EQ(expected_value,
+            *static_cast<value_type*>(TF_TensorData(actual_value.get())));
+}
+
+template <std::size_t num_devices>
+void RegisterParallelDevice(
+    TFE_Context* context, const char* device_name,
+    const std::array<const char*, num_devices>& underlying_devices,
+    TF_Status* status) {
+  TFE_CustomDevice device;
+  void* device_info;
+  tensorflow::parallel_device::AllocateParallelDevice(
+      device_name, underlying_devices.data(), underlying_devices.size(),
+      &device, &device_info);
+  TFE_RegisterCustomDevice(context, device, device_name, device_info, status);
+}
+
+}  // namespace parallel_device
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_PARALLEL_DEVICE_PARALLEL_DEVICE_TESTLIB_H_

tensorflow/c/eager/tfe_cancellation_manager_internal.h

@@ -0,0 +1,31 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_CANCELLATION_MANAGER_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_CANCELLATION_MANAGER_INTERNAL_H_
+
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/core/framework/cancellation.h"
+
+struct TFE_CancellationManager;
+typedef struct TFE_CancellationManager TFE_CancellationManager;
+
+namespace tensorflow {
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::CancellationManager,
+                            TFE_CancellationManager);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::CancellationManager*,
+                            TFE_CancellationManager*);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TFE_CANCELLATION_MANAGER_INTERNAL_H_

tensorflow/c/eager/tfe_context_internal.h

@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_CONTEXT_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_CONTEXT_INTERNAL_H_
+
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+
+// Wraps a pointer to a context implementation.
+//
+// WARNING: Since the underlying object could be ref-counted a user of this
+// interface cannot destruct the underlying context object. Instead, call
+// TFE_DeleteContext who calls Release() on the context pointer and deletes
+// the TFE_Context structure.
+typedef struct TFE_Context TFE_Context;
+
+namespace tensorflow {
+
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionContext, TFE_Context);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TFE_CONTEXT_INTERNAL_H_

tensorflow/c/eager/tfe_executor_internal.h

@@ -0,0 +1,39 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_EXECUTOR_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_EXECUTOR_INTERNAL_H_
+
+#include <memory>
+
+#include "tensorflow/core/common_runtime/eager/eager_executor.h"
+
+struct TFE_Executor {
+  explicit TFE_Executor(bool async, bool enable_streaming_enqueue,
+                        int in_flight_nodes_limit)
+      : owned_executor(new tensorflow::EagerExecutor(
+            async, enable_streaming_enqueue, in_flight_nodes_limit)) {}
+
+  explicit TFE_Executor(tensorflow::EagerExecutor* executor)
+      : owned_executor(nullptr), unowned_executor(executor) {}
+
+  tensorflow::EagerExecutor* executor() {
+    return owned_executor == nullptr ? unowned_executor : owned_executor.get();
+  }
+
+  std::unique_ptr<tensorflow::EagerExecutor> owned_executor;
+  tensorflow::EagerExecutor* unowned_executor;
+};
+
+#endif  // TENSORFLOW_C_EAGER_TFE_EXECUTOR_INTERNAL_H_

tensorflow/c/eager/tfe_monitoring_internal.h

@@ -0,0 +1,152 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_MONITORING_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_MONITORING_INTERNAL_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "absl/memory/memory.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+#include "tensorflow/core/lib/monitoring/gauge.h"
+#include "tensorflow/core/lib/monitoring/sampler.h"
+#include "tensorflow/core/platform/types.h"
+
+struct TFE_MonitoringCounterCell {
+  tensorflow::monitoring::CounterCell cell;
+};
+
+template <int NumLabels>
+struct TFE_MonitoringCounter {
+  template <typename... LabelDesc>
+  TFE_MonitoringCounter(const char* name, const char* description,
+                        LabelDesc&&... label) {
+    counter = absl::WrapUnique(tensorflow::monitoring::Counter<NumLabels>::New(
+        name, description, label...));
+  }
+
+  std::unique_ptr<tensorflow::monitoring::Counter<NumLabels>> counter;
+};
+
+struct TFE_MonitoringCounter0 : TFE_MonitoringCounter<0> {
+  using TFE_MonitoringCounter::TFE_MonitoringCounter;
+};
+struct TFE_MonitoringCounter1 : TFE_MonitoringCounter<1> {
+  using TFE_MonitoringCounter::TFE_MonitoringCounter;
+};
+struct TFE_MonitoringCounter2 : TFE_MonitoringCounter<2> {
+  using TFE_MonitoringCounter::TFE_MonitoringCounter;
+};
+
+struct TFE_MonitoringIntGaugeCell {
+  tensorflow::monitoring::GaugeCell<int64_t> cell;
+};
+struct TFE_MonitoringStringGaugeCell {
+  tensorflow::monitoring::GaugeCell<std::string> cell;
+};
+struct TFE_MonitoringBoolGaugeCell {
+  tensorflow::monitoring::GaugeCell<bool> cell;
+};
+
+template <typename ValueType, int NumLabels>
+struct TFE_MonitoringGauge {
+  template <typename... LabelDesc>
+  TFE_MonitoringGauge(const char* name, const char* description,
+                      LabelDesc&&... label) {
+    gauge = absl::WrapUnique(
+        tensorflow::monitoring::Gauge<ValueType, NumLabels>::New(
+            name, description, label...));
+  }
+
+  std::unique_ptr<tensorflow::monitoring::Gauge<ValueType, NumLabels>> gauge;
+};
+
+struct TFE_MonitoringIntGauge0 : TFE_MonitoringGauge<int64_t, 0> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringIntGauge1 : TFE_MonitoringGauge<int64_t, 1> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringIntGauge2 : TFE_MonitoringGauge<int64_t, 2> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+
+struct TFE_MonitoringStringGauge0 : TFE_MonitoringGauge<std::string, 0> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringStringGauge1 : TFE_MonitoringGauge<std::string, 1> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringStringGauge2 : TFE_MonitoringGauge<std::string, 2> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringStringGauge3 : TFE_MonitoringGauge<std::string, 3> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringStringGauge4 : TFE_MonitoringGauge<std::string, 4> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+
+struct TFE_MonitoringBoolGauge0 : TFE_MonitoringGauge<bool, 0> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringBoolGauge1 : TFE_MonitoringGauge<bool, 1> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+struct TFE_MonitoringBoolGauge2 : TFE_MonitoringGauge<bool, 2> {
+  using TFE_MonitoringGauge::TFE_MonitoringGauge;
+};
+
+struct TFE_MonitoringBuckets {
+  explicit TFE_MonitoringBuckets(
+      std::function<std::unique_ptr<tensorflow::monitoring::Buckets>(void)>
+          fn) {
+    create_buckets = fn;
+  }
+
+  std::function<std::unique_ptr<tensorflow::monitoring::Buckets>(void)>
+      create_buckets;
+};
+
+struct TFE_MonitoringSamplerCell {
+  tensorflow::monitoring::SamplerCell cell;
+};
+
+template <int NumLabels>
+struct TFE_MonitoringSampler {
+  template <typename... LabelDesc>
+  TFE_MonitoringSampler(
+      const char* name,
+      std::unique_ptr<tensorflow::monitoring::Buckets> buckets,
+      const char* description, LabelDesc&&... label) {
+    sampler = absl::WrapUnique(tensorflow::monitoring::Sampler<NumLabels>::New(
+        {name, description, label...}, std::move(buckets)));
+  }
+
+  std::unique_ptr<tensorflow::monitoring::Sampler<NumLabels>> sampler;
+};
+
+struct TFE_MonitoringSampler0 : TFE_MonitoringSampler<0> {
+  using TFE_MonitoringSampler::TFE_MonitoringSampler;
+};
+struct TFE_MonitoringSampler1 : TFE_MonitoringSampler<1> {
+  using TFE_MonitoringSampler::TFE_MonitoringSampler;
+};
+struct TFE_MonitoringSampler2 : TFE_MonitoringSampler<2> {
+  using TFE_MonitoringSampler::TFE_MonitoringSampler;
+};
+
+#endif  // TENSORFLOW_C_EAGER_TFE_MONITORING_INTERNAL_H_

tensorflow/c/eager/tfe_monitoring_reader_internal.h

@@ -0,0 +1,34 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EAGER_TFE_MONITORING_READER_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_MONITORING_READER_INTERNAL_H_
+
+#include <memory>
+
+#include "tensorflow/core/lib/monitoring/cell_reader.h"
+
+struct TFE_MonitoringCounterReader {
+  explicit TFE_MonitoringCounterReader(const char* name) {
+    counter = std::make_unique<
+        ::tensorflow::monitoring::testing::CellReader<int64_t>>(name);
+  }
+  template <typename... LabelType>
+  int64_t Read(const LabelType&... labels);
+  std::unique_ptr<::tensorflow::monitoring::testing::CellReader<int64_t>>
+      counter;
+};
+
+#endif  // TENSORFLOW_C_EAGER_TFE_MONITORING_READER_INTERNAL_H_

tensorflow/c/eager/tfe_op_attrs_internal.h

@@ -0,0 +1,39 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_OP_ATTRS_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_OP_ATTRS_INTERNAL_H_
+
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/c/eager/abstract_op_attrs.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+
+// An equivalent of a tensorflow::NameAttrList protocol buffer, but used in ways
+// that sometimes do not require serialization.
+typedef struct TFE_OpAttrs TFE_OpAttrs;
+
+typedef struct TFE_Context TFE_Context;
+typedef struct TFE_Op TFE_Op;
+
+namespace tensorflow {
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractOpAttrs, TFE_OpAttrs);
+
+// Set an AttrValue on the op. Doesn't handle the list types.
+void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op,
+                          const tensorflow::AttrValue& default_value,
+                          const char* attr_name, TF_Status* status);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TFE_OP_ATTRS_INTERNAL_H_

tensorflow/c/eager/tfe_op_internal.h

@@ -0,0 +1,36 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_OP_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_OP_INTERNAL_H_
+
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+
+// Wraps a pointer to an operation implementation.
+//
+// WARNING: Since the underlying object could be ref-counted a user of this
+// interface cannot destruct the underlying operation object. Instead, call
+// TFE_DeleteOp who calls Release() on the operation pointer and deletes
+// the TFE_Op structure.
+typedef struct TFE_Op TFE_Op;
+
+namespace tensorflow {
+
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionOperation, TFE_Op);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionOperation*, TFE_Op*);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TFE_OP_INTERNAL_H_

tensorflow/c/eager/tfe_tensor_debug_info_internal.h

@@ -0,0 +1,30 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_TENSOR_DEBUG_INFO_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_TENSOR_DEBUG_INFO_INTERNAL_H_
+
+#include <vector>
+
+#include "tensorflow/core/platform/types.h"
+
+struct TFE_TensorDebugInfo {
+  explicit TFE_TensorDebugInfo(const std::vector<int64_t>& dims)
+      : dev_dims(dims) {}
+
+  // Fully-padded, minor-to-major.
+  std::vector<int64_t> dev_dims;
+};
+
+#endif  // TENSORFLOW_C_EAGER_TFE_TENSOR_DEBUG_INFO_INTERNAL_H_

tensorflow/c/eager/tfe_tensorhandle_internal.h

@@ -0,0 +1,38 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TFE_TENSORHANDLE_INTERNAL_H_
+#define TENSORFLOW_C_EAGER_TFE_TENSORHANDLE_INTERNAL_H_
+
+#include "tensorflow/c/conversion_macros.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+
+// Wraps a pointer to a tensor handle implementation.
+//
+// WARNING: Since the underlying object could be ref-counted a user of this
+// interface cannot destruct the underlying handle object. Instead, call
+// TFE_DeleteTensorHandle who calls Release() on the handle pointer and deletes
+// the TFE_TensorHandle structure.
+typedef struct TFE_TensorHandle TFE_TensorHandle;
+
+namespace tensorflow {
+
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionTensorHandle,
+                            TFE_TensorHandle);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionTensorHandle*,
+                            TFE_TensorHandle*);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TFE_TENSORHANDLE_INTERNAL_H_

tensorflow/c/eager/tracing_utils.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/tracing_utils.h"
+
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_operation.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace tracing {
+
+absl::Status MaybeSetOpName(AbstractOperation* op, const char* op_name) {
+  if (isa<TracingOperation>(op)) {
+    TF_RETURN_IF_ERROR(dyn_cast<TracingOperation>(op)->SetOpName(op_name));
+  }
+  if (isa<gradients::TapeOperation>(op)) {
+    TF_RETURN_IF_ERROR(MaybeSetOpName(
+        dyn_cast<gradients::TapeOperation>(op)->GetBackingOperation(),
+        op_name));
+  }
+  return absl::OkStatus();
+}
+}  // namespace tracing
+}  // namespace tensorflow

tensorflow/c/eager/tracing_utils.h

@@ -0,0 +1,26 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_TRACING_UTILS_H_
+#define TENSORFLOW_C_EAGER_TRACING_UTILS_H_
+
+#include "tensorflow/c/eager/abstract_operation.h"
+
+namespace tensorflow {
+namespace tracing {
+absl::Status MaybeSetOpName(AbstractOperation*, const char* op_name);
+}  // namespace tracing
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_TRACING_UTILS_H_

tensorflow/c/eager/unified_api_test.cc

@@ -0,0 +1,198 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+class UnifiedAPI
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    absl::Status s = StatusFromTF_Status(status.get());
+    CHECK_EQ(errors::OK, s.code()) << s.message();
+  }
+
+ public:
+  bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
+  bool UseFunction() const { return std::get<2>(GetParam()); }
+};
+
+// Checks that inputs[0] is a scalar.
+absl::Status TestScalarShape(AbstractContext* ctx,
+                             absl::Span<AbstractTensorHandle* const> inputs,
+                             absl::Span<AbstractTensorHandle*> outputs) {
+  PartialTensorShape shape;
+  TF_RETURN_IF_ERROR(inputs[0]->Shape(&shape));
+  if (shape.dims() != 0) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Tensor expected to have scalar shape found rank: ", shape.dims()));
+  }
+  return absl::OkStatus();
+}
+
+TEST_P(UnifiedAPI, TestTensorShapeScalar) {
+  if (UseFunction() && UseMlir()) {
+    // TODO(b/173074167): Remove this.
+    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
+  }
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  absl::Status s = RunModel(TestScalarShape, ctx.get(),
+                            /*inputs=*/{x.get()},
+                            /*outputs=*/{},
+                            /*use_function=*/UseFunction());
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+}
+
+// Checks that inputs[0] is a matrix with shape 2x4.
+absl::Status TestTensorShape2x4(AbstractContext* ctx,
+                                absl::Span<AbstractTensorHandle* const> inputs,
+                                absl::Span<AbstractTensorHandle*> outputs) {
+  PartialTensorShape shape;
+  TF_RETURN_IF_ERROR(inputs[0]->Shape(&shape));
+  if (shape.dims() != 2) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Tensor expected to have rank 2 found rank: ", shape.dims()));
+  }
+  int64_t dim_sizes[] = {2, 4};
+  for (int i = 0; i < shape.dims(); i++) {
+    if (shape.dim_size(i) != dim_sizes[i]) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Dim ", i, " expected to be of size ", dim_sizes[i],
+                       " found: ", shape.dim_size(i)));
+    }
+  }
+  return absl::OkStatus();
+}
+
+TEST_P(UnifiedAPI, TestTensorShape2x4) {
+  if (UseFunction() && UseMlir()) {
+    // TODO(b/173074167): Remove this.
+    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
+  }
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    float data[] = {0., 0., 0., 0., 0., 0., 0., 0};
+    int64_t dim_sizes[] = {2, 4};
+    absl::Status s = TestTensorHandleWithDims<float, TF_FLOAT>(
+        ctx.get(), data, dim_sizes, 2, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  absl::Status s = RunModel(TestTensorShape2x4, ctx.get(),
+                            /*inputs=*/{x.get()},
+                            /*outputs=*/{},
+                            /*use_function=*/UseFunction());
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+}
+
+TEST_P(UnifiedAPI, TestUnknownShapeTracing) {
+  if (!UseFunction()) {
+    GTEST_SKIP() << "Tracing only test.";
+  }
+  if (UseMlir()) {
+    // TODO(b/173074167): Remove this.
+    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
+  }
+  AbstractContextPtr ctx(BuildFunction("test_fn"));
+  AbstractTensorHandlePtr x;
+  {
+    tracing::TracingTensorHandle* x_raw = nullptr;
+    PartialTensorShape shape;
+    absl::Status s = dyn_cast<tracing::TracingContext>(ctx.get())->AddParameter(
+        DT_FLOAT, shape, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  PartialTensorShape shape;
+  absl::Status s = x->Shape(&shape);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  ASSERT_TRUE(shape.unknown_rank());
+}
+
+TEST_P(UnifiedAPI, TestPartialShapeTracing) {
+  if (!UseFunction()) {
+    GTEST_SKIP() << "Tracing only test.";
+  }
+  if (UseMlir()) {
+    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
+  }
+  AbstractContextPtr ctx(BuildFunction("test_fn"));
+  AbstractTensorHandlePtr x;
+  {
+    tracing::TracingTensorHandle* x_raw = nullptr;
+    PartialTensorShape shape;
+    int64_t dim_sizes[] = {2, -1};
+    absl::Status s = PartialTensorShape::MakePartialShape(dim_sizes, 2, &shape);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    s = dyn_cast<tracing::TracingContext>(ctx.get())->AddParameter(
+        DT_FLOAT, shape, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  PartialTensorShape shape;
+  absl::Status s = x->Shape(&shape);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  ASSERT_FALSE(shape.unknown_rank());
+
+  ASSERT_EQ(2, shape.dim_size(0));
+  ASSERT_EQ(-1, shape.dim_size(1));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCppAPI, UnifiedAPI,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/eager/unified_api_testutil.cc

@@ -0,0 +1,145 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/c/eager/unified_api_testutil.h"
+
+#include "absl/container/flat_hash_set.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+
+AbstractContext* BuildFunction(const char* fn_name) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TF_ExecutionContext* graph_ctx = TF_CreateFunction(fn_name, status.get());
+  return unwrap(graph_ctx);
+}
+
+absl::Status CreateParamsForInputs(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
+    std::vector<AbstractTensorHandle*>* params) {
+  tracing::TracingTensorHandle* handle = nullptr;
+  for (auto input : inputs) {
+    PartialTensorShape shape;
+    TF_RETURN_IF_ERROR(input->Shape(&shape));
+    TF_RETURN_IF_ERROR(dyn_cast<tracing::TracingContext>(ctx)->AddParameter(
+        input->DataType(), shape, &handle));
+    params->emplace_back(handle);
+  }
+  return absl::OkStatus();
+}
+
+// Runs `model` maybe wrapped in a function.
+absl::Status RunModel(Model model, AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs,
+                      bool use_function) {
+  if (use_function) {
+    const char* fn_name = "test_fn";
+    core::RefCountPtr<AbstractFunction> scoped_func;
+    // Returning null tensors from a tf.function is not supported, so we keep
+    // track of indices in the model's outputs are nullptr in this set.
+    // The FunctionDef only outputs the non-null tensors. We later pad the
+    // function op outputs to have nullptrs at the `null_indices`.
+    absl::flat_hash_set<int> null_indices;
+    {
+      AbstractContextPtr func_ctx(BuildFunction(fn_name));
+      std::vector<AbstractTensorHandle*> func_inputs;
+      func_inputs.reserve(inputs.size());
+      TF_RETURN_IF_ERROR(
+          CreateParamsForInputs(func_ctx.get(), inputs, &func_inputs));
+      std::vector<AbstractTensorHandle*> model_outputs;
+      model_outputs.resize(outputs.size());
+      TF_RETURN_IF_ERROR(model(func_ctx.get(), absl::MakeSpan(func_inputs),
+                               absl::MakeSpan(model_outputs)));
+      for (auto func_input : func_inputs) {
+        func_input->Unref();
+      }
+      AbstractFunction* func = nullptr;
+      OutputList output_list;
+      output_list.expected_num_outputs = 0;
+      output_list.outputs.reserve(outputs.size());
+      for (int i = 0; i < model_outputs.size(); i++) {
+        if (model_outputs[i]) {
+          output_list.outputs.emplace_back(model_outputs[i]);
+          output_list.expected_num_outputs += 1;
+        } else {
+          null_indices.insert(i);
+        }
+      }
+      TF_RETURN_IF_ERROR(dyn_cast<tracing::TracingContext>(func_ctx.get())
+                             ->Finalize(&output_list, &func));
+      scoped_func.reset(func);
+      for (auto output : output_list.outputs) {
+        output->Unref();
+      }
+      TF_RETURN_IF_ERROR(ctx->RegisterFunction(func));
+    }
+
+    AbstractOperationPtr fn_op(ctx->CreateOperation());
+    TF_RETURN_IF_ERROR(fn_op->Reset(fn_name, /*raw_device_name=*/nullptr));
+    for (auto input : inputs) {
+      TF_RETURN_IF_ERROR(fn_op->AddInput(input));
+    }
+    int retvals = outputs.size() - null_indices.size();
+    std::vector<AbstractTensorHandle*> fn_outputs(retvals);
+    TF_RETURN_IF_ERROR(fn_op->Execute(
+        absl::Span<AbstractTensorHandle*>(fn_outputs.data(), fn_outputs.size()),
+        &retvals));
+    int skipped_indices = 0;
+    for (int i = 0; i < outputs.size(); i++) {
+      if (!null_indices.contains(i)) {
+        outputs[i] = fn_outputs[i - skipped_indices];
+      } else {
+        skipped_indices += 1;
+      }
+    }
+    TF_RETURN_IF_ERROR(ctx->RemoveFunction(fn_name));
+    return absl::OkStatus();
+  } else {
+    return model(ctx, inputs, outputs);
+  }
+}
+
+absl::Status BuildImmediateExecutionContext(bool use_tfrt,
+                                            AbstractContext** ctx) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_ContextOptions* opts = TFE_NewContextOptions();
+  TFE_ContextOptionsSetTfrt(opts, use_tfrt);
+  *ctx = unwrap(TF_NewEagerExecutionContext(opts, status.get()));
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(status.get()));
+  TFE_DeleteContextOptions(opts);
+  return absl::OkStatus();
+}
+
+absl::Status GetValue(AbstractTensorHandle* t, TF_Tensor** result_tensor) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_TensorHandle* result_t =
+      TF_AbstractTensorGetEagerTensor(wrap(t), status.get());
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(status.get()));
+  *result_tensor = TFE_TensorHandleResolve(result_t, status.get());
+  return StatusFromTF_Status(status.get());
+}
+
+}  // namespace tensorflow

tensorflow/c/eager/unified_api_testutil.h

@@ -0,0 +1,95 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_UNIFIED_API_TESTUTIL_H_
+#define TENSORFLOW_C_EAGER_UNIFIED_API_TESTUTIL_H_
+
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// Builds and returns a `TracingContext` using the default tracing impl.
+AbstractContext* BuildFunction(const char* fn_name);
+
+// Creates parameters (placeholders) in the tracing `ctx` using the shape and
+// dtype of `inputs`.
+absl::Status CreateParamsForInputs(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
+    std::vector<AbstractTensorHandle*>* params);
+
+// A callable that takes tensor inputs and returns zero or more tensor outputs.
+using Model = std::function<absl::Status(
+    AbstractContext*, absl::Span<AbstractTensorHandle* const>,
+    absl::Span<AbstractTensorHandle*>)>;
+
+// Runs `model` maybe wrapped in a function call op. This can be thought as
+// being equivalent to the following python code.
+//
+// if use_function:
+//   outputs = tf.function(model)(inputs)
+// else:
+//   outputs = model(inputs)
+absl::Status RunModel(Model model, AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs,
+                      bool use_function);
+
+absl::Status BuildImmediateExecutionContext(bool use_tfrt,
+                                            AbstractContext** ctx);
+
+// Return a tensor handle with given type, values and dimensions.
+template <class T, TF_DataType datatype>
+absl::Status TestTensorHandleWithDims(AbstractContext* ctx, const T* data,
+                                      const int64_t* dims, int num_dims,
+                                      AbstractTensorHandle** tensor) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_Context* eager_ctx =
+      TF_ExecutionContextGetTFEContext(wrap(ctx), status.get());
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(status.get()));
+  TFE_TensorHandle* input_eager =
+      TestTensorHandleWithDims<T, datatype>(eager_ctx, data, dims, num_dims);
+  *tensor =
+      unwrap(TF_CreateAbstractTensorFromEagerTensor(input_eager, status.get()));
+  return absl::OkStatus();
+}
+
+// Return a scalar tensor handle with given value.
+template <class T, TF_DataType datatype>
+absl::Status TestScalarTensorHandle(AbstractContext* ctx, const T value,
+                                    AbstractTensorHandle** tensor) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  TFE_Context* eager_ctx =
+      TF_ExecutionContextGetTFEContext(wrap(ctx), status.get());
+  TF_RETURN_IF_ERROR(StatusFromTF_Status(status.get()));
+  TFE_TensorHandle* input_eager =
+      TestScalarTensorHandle<T, datatype>(eager_ctx, value);
+  *tensor =
+      unwrap(TF_CreateAbstractTensorFromEagerTensor(input_eager, status.get()));
+  return absl::OkStatus();
+}
+
+// Places data from `t` into *result_tensor.
+absl::Status GetValue(AbstractTensorHandle* t, TF_Tensor** result_tensor);
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_UNIFIED_API_TESTUTIL_H_

tensorflow/c/env.cc

@@ -0,0 +1,219 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/env.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/tf_file_statistics.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_statistics.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+
+struct TF_StringStream {
+  std::vector<std::string>* list;
+  size_t position;
+};
+
+void TF_CreateDir(const char* dirname, TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->CreateDir(dirname));
+}
+
+void TF_DeleteDir(const char* dirname, TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->DeleteDir(dirname));
+}
+
+void TF_DeleteRecursively(const char* dirname, int64_t* undeleted_file_count,
+                          int64_t* undeleted_dir_count, TF_Status* status) {
+  ::int64_t f, d;
+
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->DeleteRecursively(dirname, &f, &d));
+  *undeleted_file_count = f;
+  *undeleted_dir_count = d;
+}
+
+void TF_FileStat(const char* filename, TF_FileStatistics* stats,
+                 TF_Status* status) {
+  ::tensorflow::FileStatistics cc_stats;
+  TF_SetStatus(status, TF_OK, "");
+  absl::Status s = ::tensorflow::Env::Default()->Stat(filename, &cc_stats);
+  ::tensorflow::Set_TF_Status_from_Status(status, s);
+  if (s.ok()) {
+    stats->length = cc_stats.length;
+    stats->mtime_nsec = cc_stats.mtime_nsec;
+    stats->is_directory = cc_stats.is_directory;
+  }
+}
+
+void TF_NewWritableFile(const char* filename, TF_WritableFileHandle** handle,
+                        TF_Status* status) {
+  std::unique_ptr<::tensorflow::WritableFile> f;
+  TF_SetStatus(status, TF_OK, "");
+  absl::Status s = ::tensorflow::Env::Default()->NewWritableFile(filename, &f);
+  ::tensorflow::Set_TF_Status_from_Status(status, s);
+
+  if (s.ok()) {
+    *handle = reinterpret_cast<TF_WritableFileHandle*>(f.release());
+  }
+}
+
+void TF_CloseWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+  auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Close());
+  delete cc_file;
+}
+
+void TF_SyncWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+  auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Sync());
+}
+
+void TF_FlushWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+  auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Flush());
+}
+
+void TF_AppendWritableFile(TF_WritableFileHandle* handle, const char* data,
+                           size_t length, TF_Status* status) {
+  auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, cc_file->Append(absl::string_view{data, length}));
+}
+
+void TF_DeleteFile(const char* filename, TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->DeleteFile(filename));
+}
+
+bool TF_StringStreamNext(TF_StringStream* list, const char** result) {
+  if (list->position >= list->list->size()) {
+    *result = nullptr;
+    return false;
+  }
+
+  *result = list->list->at(list->position++).c_str();
+  return true;
+}
+
+void TF_StringStreamDone(TF_StringStream* list) {
+  delete list->list;
+  delete list;
+}
+TF_StringStream* TF_GetChildren(const char* dirname, TF_Status* status) {
+  auto* children = new std::vector<std::string>;
+
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->GetChildren(dirname, children));
+
+  auto* list = new TF_StringStream;
+  list->list = children;
+  list->position = 0;
+  return list;
+}
+
+TF_StringStream* TF_GetLocalTempDirectories() {
+  auto* tmpdirs = new std::vector<std::string>;
+
+  ::tensorflow::Env::Default()->GetLocalTempDirectories(tmpdirs);
+
+  auto* list = new TF_StringStream;
+  list->list = tmpdirs;
+  list->position = 0;
+  return list;
+}
+
+char* TF_GetTempFileName(const char* extension) {
+  return strdup(::tensorflow::io::GetTempFilename(extension).c_str());
+}
+
+TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) {
+  return ::tensorflow::Env::Default()->NowNanos();
+}
+
+// Returns the number of microseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void) {
+  return ::tensorflow::Env::Default()->NowMicros();
+}
+
+// Returns the number of seconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void) {
+  return ::tensorflow::Env::Default()->NowSeconds();
+}
+
+void TF_DefaultThreadOptions(TF_ThreadOptions* options) {
+  options->stack_size = 0;
+  options->guard_size = 0;
+  options->numa_node = -1;
+}
+
+TF_Thread* TF_StartThread(const TF_ThreadOptions* options,
+                          const char* thread_name, void (*work_func)(void*),
+                          void* param) {
+  ::tensorflow::ThreadOptions cc_options;
+  cc_options.stack_size = options->stack_size;
+  cc_options.guard_size = options->guard_size;
+  cc_options.numa_node = options->numa_node;
+  return reinterpret_cast<TF_Thread*>(::tensorflow::Env::Default()->StartThread(
+      cc_options, thread_name, [=]() { (*work_func)(param); }));
+}
+
+void TF_JoinThread(TF_Thread* thread) {
+  // ::tensorflow::Thread joins on destruction
+  delete reinterpret_cast<::tensorflow::Thread*>(thread);
+}
+
+void* TF_LoadSharedLibrary(const char* library_filename, TF_Status* status) {
+  void* handle = nullptr;
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->LoadDynamicLibrary(library_filename,
+                                                               &handle));
+  return handle;
+}
+
+void* TF_GetSymbolFromLibrary(void* handle, const char* symbol_name,
+                              TF_Status* status) {
+  void* symbol = nullptr;
+  TF_SetStatus(status, TF_OK, "");
+  ::tensorflow::Set_TF_Status_from_Status(
+      status, ::tensorflow::Env::Default()->GetSymbolFromLibrary(
+                  handle, symbol_name, &symbol));
+  return symbol;
+}

tensorflow/c/env.h

@@ -0,0 +1,212 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_ENV_H_
+#define TENSORFLOW_C_ENV_H_
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/tf_file_statistics.h"
+#include "tensorflow/c/tf_status.h"
+
+// --------------------------------------------------------------------------
+// C API for tensorflow::Env.
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct TF_WritableFileHandle TF_WritableFileHandle;
+typedef struct TF_StringStream TF_StringStream;
+typedef struct TF_Thread TF_Thread;
+
+typedef struct TF_ThreadOptions {
+  // Thread stack size to use (in bytes), zero implies that the system default
+  // will be used.
+  size_t stack_size;
+
+  // Guard area size to use near thread stacks to use (in bytes), zero implies
+  // that the system default will be used.
+  size_t guard_size;
+
+  // The NUMA node to use, -1 implies that there should be no NUMA affinity for
+  // this thread.
+  int numa_node;
+} TF_ThreadOptions;
+
+// Creates the specified directory. Typical status code are:
+//  * TF_OK - successfully created the directory
+//  * TF_ALREADY_EXISTS - directory already exists
+//  * TF_PERMISSION_DENIED - dirname is not writable
+TF_CAPI_EXPORT extern void TF_CreateDir(const char* dirname, TF_Status* status);
+
+// Deletes the specified directory. Typical status codes are:
+//  * TF_OK - successfully deleted the directory
+//  * TF_FAILED_PRECONDITION - the directory is not empty
+TF_CAPI_EXPORT extern void TF_DeleteDir(const char* dirname, TF_Status* status);
+
+// Deletes the specified directory and all subdirectories and files underneath
+// it. This is accomplished by traversing the directory tree rooted at dirname
+// and deleting entries as they are encountered.
+//
+// If dirname itself is not readable or does not exist, *undeleted_dir_count is
+// set to 1, *undeleted_file_count is set to 0 and an appropriate status (e.g.
+// TF_NOT_FOUND) is returned.
+//
+// If dirname and all its descendants were successfully deleted, TF_OK is
+// returned and both error counters are set to zero.
+//
+// Otherwise, while traversing the tree, undeleted_file_count and
+// undeleted_dir_count are updated if an entry of the corresponding type could
+// not be deleted. The returned error status represents the reason that any one
+// of these entries could not be deleted.
+//
+// Typical status codes:
+//  * TF_OK - dirname exists and we were able to delete everything underneath
+//  * TF_NOT_FOUND - dirname doesn't exist
+//  * TF_PERMISSION_DENIED - dirname or some descendant is not writable
+//  * TF_UNIMPLEMENTED - some underlying functions (like Delete) are not
+//    implemented
+TF_CAPI_EXPORT extern void TF_DeleteRecursively(const char* dirname,
+                                                int64_t* undeleted_file_count,
+                                                int64_t* undeleted_dir_count,
+                                                TF_Status* status);
+
+// Obtains statistics for the given path. If status is TF_OK, *stats is
+// updated, otherwise it is not touched.
+TF_CAPI_EXPORT extern void TF_FileStat(const char* filename,
+                                       TF_FileStatistics* stats,
+                                       TF_Status* status);
+
+// Creates or truncates the given filename and returns a handle to be used for
+// appending data to the file. If status is TF_OK, *handle is updated and the
+// caller is responsible for freeing it (see TF_CloseWritableFile).
+TF_CAPI_EXPORT extern void TF_NewWritableFile(const char* filename,
+                                              TF_WritableFileHandle** handle,
+                                              TF_Status* status);
+
+// Closes the given handle and frees its memory. If there was a problem closing
+// the file, it is indicated by status. Memory is freed in any case.
+TF_CAPI_EXPORT extern void TF_CloseWritableFile(TF_WritableFileHandle* handle,
+                                                TF_Status* status);
+
+// Syncs content of the handle to the filesystem. Blocks waiting for the
+// filesystem to indicate that the content has been persisted.
+TF_CAPI_EXPORT extern void TF_SyncWritableFile(TF_WritableFileHandle* handle,
+                                               TF_Status* status);
+
+// Flush local buffers to the filesystem. If the process terminates after a
+// successful flush, the contents may still be persisted, since the underlying
+// filesystem may eventually flush the contents.  If the OS or machine crashes
+// after a successful flush, the contents may or may not be persisted, depending
+// on the implementation.
+TF_CAPI_EXPORT extern void TF_FlushWritableFile(TF_WritableFileHandle* handle,
+                                                TF_Status* status);
+
+// Appends the given bytes to the file. Any failure to do so is indicated in
+// status.
+TF_CAPI_EXPORT extern void TF_AppendWritableFile(TF_WritableFileHandle* handle,
+                                                 const char* data,
+                                                 size_t length,
+                                                 TF_Status* status);
+
+// Deletes the named file and indicates whether successful in *status.
+TF_CAPI_EXPORT extern void TF_DeleteFile(const char* filename,
+                                         TF_Status* status);
+
+// Retrieves the next item from the given TF_StringStream and places a pointer
+// to it in *result. If no more items are in the list, *result is set to NULL
+// and false is returned.
+//
+// Ownership of the items retrieved with this function remains with the library.
+// Item points are invalidated after a call to TF_StringStreamDone.
+TF_CAPI_EXPORT extern bool TF_StringStreamNext(TF_StringStream* list,
+                                               const char** result);
+
+// Frees the resources associated with given string list. All pointers returned
+// by TF_StringStreamNext are invalid after this call.
+TF_CAPI_EXPORT extern void TF_StringStreamDone(TF_StringStream* list);
+
+// Retrieves the list of children of the given directory. You can iterate
+// through the list with TF_StringStreamNext. The caller is responsible for
+// freeing the list (see TF_StringStreamDone).
+TF_CAPI_EXPORT extern TF_StringStream* TF_GetChildren(const char* filename,
+                                                      TF_Status* status);
+
+// Retrieves a list of directory names on the local machine that may be used for
+// temporary storage. You can iterate through the list with TF_StringStreamNext.
+// The caller is responsible for freeing the list (see TF_StringStreamDone).
+TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void);
+
+// Creates a temporary file name with an extension.
+// The caller is responsible for freeing the returned pointer.
+TF_CAPI_EXPORT extern char* TF_GetTempFileName(const char* extension);
+
+// Returns the number of nanoseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void);
+
+// Returns the number of microseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void);
+
+// Returns the number of seconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void);
+
+// Populates a TF_ThreadOptions struct with system-default values.
+TF_CAPI_EXPORT extern void TF_DefaultThreadOptions(TF_ThreadOptions* options);
+
+// Returns a new thread that is running work_func and is identified
+// (for debugging/performance-analysis) by thread_name.
+//
+// The given param (which may be null) is passed to work_func when the thread
+// starts. In this way, data may be passed from the thread back to the caller.
+//
+// Caller takes ownership of the result and must call TF_JoinThread on it
+// eventually.
+TF_CAPI_EXPORT extern TF_Thread* TF_StartThread(const TF_ThreadOptions* options,
+                                                const char* thread_name,
+                                                void (*work_func)(void*),
+                                                void* param);
+
+// Waits for the given thread to finish execution, then deletes it.
+TF_CAPI_EXPORT extern void TF_JoinThread(TF_Thread* thread);
+
+// \brief Load a dynamic library.
+//
+// Pass "library_filename" to a platform-specific mechanism for dynamically
+// loading a library. The rules for determining the exact location of the
+// library are platform-specific and are not documented here.
+//
+// On success, place OK in status and return the newly created library handle.
+// Otherwise returns nullptr and set error status.
+TF_CAPI_EXPORT extern void* TF_LoadSharedLibrary(const char* library_filename,
+                                                 TF_Status* status);
+
+// \brief Get a pointer to a symbol from a dynamic library.
+//
+// "handle" should be a pointer returned from a previous call to
+// TF_LoadLibraryFromEnv. On success, place OK in status and return a pointer to
+// the located symbol. Otherwise returns nullptr and set error status.
+TF_CAPI_EXPORT extern void* TF_GetSymbolFromLibrary(void* handle,
+                                                    const char* symbol_name,
+                                                    TF_Status* status);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // TENSORFLOW_C_ENV_H_