~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,

  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;

  const int len = std::min(num_values, static_cast<int>(nodes.size()));
  char* p = static_cast<char*>(storage);
  for (int index = 0; index < len; ++index) {
    const std::string& s = nodes[index];
    values[index] = p;
    lengths[index] = s.size();
    if ((p + s.size()) > (static_cast<char*>(storage) + storage_size)) {
      tsl::Set_TF_Status_from_Status(
          status,

    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) =
      tensorflow::strings::StrCat("localhost:", port);

    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) {
  Status s;
  AbstractTensorHandlePtr numerical_grad;
  {
    AbstractTensorHandle* numerical_grad_raw;
    s = CalcNumericalGrad(ctx, model, inputs, input_index, use_function,
                          &numerical_grad_raw);

  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(

  // 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(

    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];
}

  int num_inputs = TFE_OpGetFlatInputCount(other, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  for (int input_index = 0; input_index < num_inputs; ++input_index) {
    TFE_TensorHandle* input = TFE_OpGetFlatInput(other, input_index, status);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    TFE_OpAddInput(ret, input, status);

  return num_elements;
}

int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index,
                            TF_Status* status) {
  if (h == nullptr) {
    status->status = tensorflow::errors::InvalidArgument("Invalid handle");
    return -1;
  }

  int64_t dim = -1;
  status->status = tensorflow::unwrap(h)->Dim(dim_index, &dim);
  return dim;
}