| 70 | } |
| 71 | |
| 72 | TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { |
| 73 | const TfLiteTensor* axis = GetInput(context, node, 0); |
| 74 | const TfLiteTensor* input = GetInput(context, node, 1); |
| 75 | |
| 76 | // Dynamic output tensors are needed if axis tensor is not constant. |
| 77 | // But Micro doesn't support dynamic memeory allocation, so we only support |
| 78 | // constant axis tensor for now. |
| 79 | TF_LITE_ENSURE_MSG(context, IsConstantTensor(axis), |
| 80 | "Non constant axis tensor not supported"); |
| 81 | |
| 82 | int axis_value = GetTensorData<int32_t>(axis)[0]; |
| 83 | if (axis_value < 0) { |
| 84 | axis_value += NumDimensions(input); |
| 85 | } |
| 86 | |
| 87 | TF_LITE_ENSURE(context, axis_value >= 0); |
| 88 | TF_LITE_ENSURE(context, axis_value < NumDimensions(input)); |
| 89 | |
| 90 | switch (input->type) { |
| 91 | case kTfLiteFloat32: { |
| 92 | return SplitImpl<float>(context, node, input, axis_value); |
| 93 | } |
| 94 | case kTfLiteUInt8: { |
| 95 | return SplitImpl<uint8_t>(context, node, input, axis_value); |
| 96 | } |
| 97 | case kTfLiteInt8: { |
| 98 | return SplitImpl<int8_t>(context, node, input, axis_value); |
| 99 | } |
| 100 | case kTfLiteInt16: { |
| 101 | return SplitImpl<int16_t>(context, node, input, axis_value); |
| 102 | } |
| 103 | case kTfLiteInt32: { |
| 104 | return SplitImpl<int32_t>(context, node, input, axis_value); |
| 105 | } |
| 106 | default: |
| 107 | context->ReportError(context, "Type %s currently not supported.", |
| 108 | TfLiteTypeGetName(input->type)); |
| 109 | return kTfLiteError; |
| 110 | } |
| 111 | #undef TF_LITE_SPLIT |
| 112 | |
| 113 | return kTfLiteOk; |
| 114 | } |
| 115 | |
| 116 | } // namespace split |
| 117 |
nothing calls this directly
no test coverage detected