| 186 | } // namespace |
| 187 | |
| 188 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 189 | // Check that the inputs and outputs have the right sizes and types. |
| 190 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); |
| 191 | TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2); |
| 192 | |
| 193 | const TfLiteTensor* input = GetInput(context, node, kInputTensor); |
| 194 | TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); |
| 195 | TF_LITE_ENSURE_EQ(context, input->type, output_values->type); |
| 196 | |
| 197 | const TfLiteTensor* top_k = GetInput(context, node, kInputTopK); |
| 198 | TF_LITE_ENSURE_EQ(context, top_k->type, kTfLiteInt32); |
| 199 | |
| 200 | // Set output dynamic if the input is not const. |
| 201 | if (IsConstantTensor(top_k)) { |
| 202 | TF_LITE_ENSURE_OK(context, ResizeOutput(context, node)); |
| 203 | } else { |
| 204 | TfLiteTensor* output_indexes = GetOutput(context, node, kOutputIndexes); |
| 205 | TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); |
| 206 | SetTensorToDynamic(output_indexes); |
| 207 | SetTensorToDynamic(output_values); |
| 208 | } |
| 209 | return kTfLiteOk; |
| 210 | } |
| 211 | |
| 212 | TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { |
| 213 | TfLiteTensor* output_values = GetOutput(context, node, kOutputValues); |
nothing calls this directly
no test coverage detected