| 30 | constexpr int kOutputTensor = 0; |
| 31 | |
| 32 | TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input, |
| 33 | const TfLiteTensor* axis, TfLiteTensor* output) { |
| 34 | int axis_value = *GetTensorData<int>(axis); |
| 35 | if (axis_value < 0) { |
| 36 | axis_value += NumDimensions(input); |
| 37 | } |
| 38 | |
| 39 | // Copy the input dimensions to output except the axis dimension. |
| 40 | TfLiteIntArray* output_dims = TfLiteIntArrayCreate(NumDimensions(input) - 1); |
| 41 | int j = 0; |
| 42 | for (int i = 0; i < NumDimensions(input); ++i) { |
| 43 | if (i != axis_value) { |
| 44 | output_dims->data[j] = SizeOfDimension(input, i); |
| 45 | ++j; |
| 46 | } |
| 47 | } |
| 48 | return context->ResizeTensor(context, output, output_dims); |
| 49 | } |
| 50 | |
| 51 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 52 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); |
no test coverage detected