| 30 | constexpr int kOutputTensor = 0; |
| 31 | |
| 32 | TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) { |
| 33 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 3); |
| 34 | TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); |
| 35 | |
| 36 | const TfLiteTensor* input_condition = |
| 37 | GetInput(context, node, kInputTensorCondition); |
| 38 | const TfLiteTensor* input_x = GetInput(context, node, kInputTensorX); |
| 39 | const TfLiteTensor* input_y = GetInput(context, node, kInputTensorY); |
| 40 | TfLiteTensor* output = GetOutput(context, node, kOutputTensor); |
| 41 | |
| 42 | // Input must be bool. |
| 43 | TF_LITE_ENSURE(context, input_condition->type == kTfLiteBool); |
| 44 | |
| 45 | // Input tensors must have the same type and size |
| 46 | TF_LITE_ENSURE_EQ(context, input_x->type, input_y->type); |
| 47 | TF_LITE_ENSURE(context, HaveSameShapes(input_x, input_y)); |
| 48 | output->type = input_x->type; |
| 49 | |
| 50 | // Either the same shape, or input_condition must be Rank 1 and match over the |
| 51 | // first dimension. |
| 52 | bool same_shape = HaveSameShapes(input_condition, input_x); |
| 53 | if (!same_shape && NumDimensions(input_condition) == 1) { |
| 54 | same_shape = |
| 55 | SizeOfDimension(input_condition, 0) == SizeOfDimension(input_x, 0); |
| 56 | } |
| 57 | |
| 58 | TF_LITE_ENSURE(context, same_shape); |
| 59 | |
| 60 | TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_x->dims); |
| 61 | return context->ResizeTensor(context, output, output_size); |
| 62 | } |
| 63 | |
| 64 | TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) { |
| 65 | const TfLiteTensor* input_condition = |
nothing calls this directly
no test coverage detected