const int64_t* dims, int num_dims) {
  DCHECK(index >= 0 && index < shape_list->num_items);
  TF_ShapeAndType& shape = shape_list->items[index];
  DCHECK(shape.dims == nullptr) << "Shape at " << index << " is already set!";
  DCHECK(num_dims >= 0) << "Number of dimensions cannot be negative!";
  shape.num_dims = num_dims;
  shape.dims = new int64_t[num_dims];

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

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"
      "      }"

  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,

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

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

                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    // TODO(b/161805092): Support broadcasting.

    DCHECK(grad_outputs[0]);
    grad_inputs[0] = grad_outputs[0];
    grad_inputs[1] = grad_outputs[0];

    grad_inputs[0]->Ref();
    grad_inputs[1]->Ref();
    return absl::OkStatus();
  }

  }

  tensorflow::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();
    }

  TF_VLog(1, "Successful read of %s @ %u of size: %u", path.c_str(), offset,
          read);
  stream.read(buffer, read);
  read = stream.gcount();
  if (read < buffer_size) {
    // Check stat cache to see if we encountered an interrupted read.
    tf_gcs_filesystem::GcsFileStat stat;
    if (gcs_file->stat_cache->Lookup(path, &stat)) {
      if (offset + read < stat.base.length) {

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