MyComputeFunc(kernel, ctx);

    TF_Status* status = TF_NewStatus();

    TF_SetStatus(status, TF_OK, "");
    EXPECT_EQ(false, TF_IsHostMemoryInput(ctx, 0, status));
    EXPECT_EQ(TF_OK, TF_GetCode(status));

    TF_SetStatus(status, TF_OK, "");
    EXPECT_EQ(true, TF_IsHostMemoryInput(ctx, 1, status));
    EXPECT_EQ(TF_OK, TF_GetCode(status));

    TF_SetStatus(status, TF_OK, "");

    TF_OpKernelContext* ctx, int i);

// Returns true if the ith input is allocated in host memory. If i < 0 or i >=
// TF_NumInputs(ctx), the program aborts.
TF_CAPI_EXPORT extern bool TF_IsHostMemoryInput(TF_OpKernelContext* ctx, int i,
                                                TF_Status* status);

// Returns true if the ith output is allocated in host memory. If i < 0 or i >=

  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  CHECK_GE(i, 0);
  CHECK_LT(i, cc_ctx->num_outputs());
  return static_cast<TF_DataType>(cc_ctx->expected_output_dtype(i));
}

bool TF_IsHostMemoryInput(TF_OpKernelContext* ctx, int i, TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  if (i < 0 || i >= cc_ctx->num_inputs()) {