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

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

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

  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;

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

      << "], 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);

  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>

                         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(