| 28 | constexpr int kOutputTensor = 0; |
| 29 | |
| 30 | TfLiteStatus ResizeOutputTensor(TfLiteContext* context, |
| 31 | const TfLiteTensor* cond_tensor, |
| 32 | TfLiteTensor* output_tensor) { |
| 33 | // Output tensor should have shape: |
| 34 | // (num_true, cond_rank), where num_true denotes the number of true values |
| 35 | // in condition. |
| 36 | const RuntimeShape& cond_shape = GetTensorShape(cond_tensor); |
| 37 | const int size = cond_shape.FlatSize(); |
| 38 | const int cond_rank = cond_shape.DimensionsCount(); |
| 39 | const bool* cond_data = GetTensorData<bool>(cond_tensor); |
| 40 | |
| 41 | int true_count = 0; |
| 42 | for (int i = 0; i < size; ++i) { |
| 43 | if (cond_data[i]) { |
| 44 | true_count++; |
| 45 | } |
| 46 | } |
| 47 | TfLiteIntArray* output_dims = TfLiteIntArrayCreate(2); |
| 48 | output_dims->data[0] = true_count; |
| 49 | output_dims->data[1] = cond_rank; |
| 50 | return context->ResizeTensor(context, output_tensor, output_dims); |
| 51 | } |
| 52 | |
| 53 | TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { |
| 54 | TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); |
no test coverage detected