if (i < 0 || i >= cc_ctx->num_inputs()) {
    TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range");
    return false;
  }
  TF_SetStatus(status, TF_OK, "");
  return cc_ctx->input_memory_type(i) == tensorflow::HOST_MEMORY;
}

bool TF_IsHostMemoryOutput(TF_OpKernelContext* ctx, int i, TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);

  TF_RETURN_IF_ERROR(GetCompileTimeConstInputs(
      &ctx->op_kernel(), &constant_input_indices, ctx->function_library()));
  if (!absl::c_all_of(constant_input_indices, [&](int idx) {
        return ctx->input_memory_type(idx) == HOST_MEMORY;
      })) {
    return errors::Internal("Unexpected device placement for a constant input");
  }
  return constant_input_indices;
}