std::vector<int> resource_arg_indices;
  TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
      flr, function, &fbody, &constant_arg_indices, &resource_arg_indices));

  MemoryTypeVector input_memory_types =
      GetInputMemoryTypes(fbody, constant_arg_indices, resource_arg_indices);
  MemoryTypeVector output_memory_types = GetOutputMemoryTypes(fbody);

  // Create the kernel.
  Device* dev = flr->device();

  EXPECT_EQ(2, kernel_->num_inputs());
  EXPECT_EQ(DT_FLOAT, kernel_->input_type(0));
  EXPECT_EQ(DT_RESOURCE, kernel_->input_type(1));
  EXPECT_EQ(DEVICE_MEMORY, kernel_->input_memory_types()[0]);
  EXPECT_EQ(HOST_MEMORY, kernel_->input_memory_types()[1]);

  EXPECT_EQ(1, kernel_->num_outputs());
  EXPECT_EQ(DT_FLOAT, kernel_->output_type(0));
  EXPECT_EQ(DEVICE_MEMORY, kernel_->output_memory_types()[0]);
}

    if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
      // Compile-time constants and resource handles are expected to be in
      // host memory.
      input_memory_types[i] = tensorflow::HOST_MEMORY;
    }
  }
  return input_memory_types;
}

tensorflow::MemoryTypeVector GetOutputMemoryTypes(
    const tensorflow::FunctionBody* fbody) {

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