| 102 | } |
| 103 | |
| 104 | TfLiteStatus ResizeOutputTensor(TfLiteContext* context, |
| 105 | const OneHotContext& op_context) { |
| 106 | TF_LITE_ENSURE(context, *op_context.depth->data.i32 >= 0); |
| 107 | TfLiteIntArray* output_size = TfLiteIntArrayCreate(op_context.output_dims); |
| 108 | for (int i = 0; i < op_context.output_dims; ++i) { |
| 109 | if (i < op_context.axis) { |
| 110 | output_size->data[i] = op_context.indices->dims->data[i]; |
| 111 | } else if (i == op_context.axis) { |
| 112 | output_size->data[i] = *op_context.depth->data.i32; |
| 113 | } else { |
| 114 | output_size->data[i] = op_context.indices->dims->data[i - 1]; |
| 115 | } |
| 116 | } |
| 117 | return context->ResizeTensor(context, op_context.output, output_size); |
| 118 | } |
| 119 | |
| 120 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 121 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 4); |
no test coverage detected