TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  // Check that the inputs and outputs have the right sizes and types.
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);

  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  TfLiteTensor* output_values;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, kOutputValues, &output_values));
  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output_values->type);

  const TfLiteTensor* top_k;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTopK, &top_k));
  TF_LITE_ENSURE_TYPES_EQ(context, top_k->type, kTfLiteInt32);

  // Set output dynamic if the input is not const.
  if (IsConstantTensor(top_k)) {
    TF_LITE_ENSURE_OK(context, ResizeOutput(context, node));
  } else {
    TfLiteTensor* output_indexes;
    TF_LITE_ENSURE_OK(
        context, GetOutputSafe(context, node, kOutputIndexes, &output_indexes));
    TfLiteTensor* output_values;
    TF_LITE_ENSURE_OK(
        context, GetOutputSafe(context, node, kOutputValues, &output_values));
    SetTensorToDynamic(output_indexes);
    SetTensorToDynamic(output_values);
  }
  return kTfLiteOk;
}
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);

  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  TfLiteTensor* output_values;
  TF_LITE_ENSURE_OK(
      context, GetOutputSafe(context, node, kOutputValues, &output_values));
  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output_values->type);

  const TfLiteTensor* top_k;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTopK, &top_k));
  TF_LITE_ENSURE_TYPES_EQ(context, top_k->type, kTfLiteInt32);

   
  if (IsConstantTensor(top_k)) {
    TF_LITE_ENSURE_OK(context, ResizeOutput(context, node));
  } else {
    TfLiteTensor* output_indexes;
    TF_LITE_ENSURE_OK(
        context, GetOutputSafe(context, node, kOutputIndexes, &output_indexes));
    TfLiteTensor* output_values;
    TF_LITE_ENSURE_OK(
        context, GetOutputSafe(context, node, kOutputValues, &output_values));
    SetTensorToDynamic(output_indexes);
    SetTensorToDynamic(output_values);
  }
  return kTfLiteOk;
}